aerospike 2.22.0 → 2.24.0

data/lib/aerospike/query/partition_tracker.rb

@@ -0,0 +1,347 @@
+# frozen_string_literal: true
+
+# Copyright 2014-2020 Aerospike, Inc.
+#
+# Portions may be licensed to Aerospike, Inc. under one or more contributor
+# license agreements.
+#
+# Licensed 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.
+
+module Aerospike
+  class PartitionTracker
+    attr_reader :partitions, :partitions_capacity, :partition_begin, :node_capacity,
+                :node_filter, :partition_filter, :node_partitions_list, :max_records,
+                :sleep_between_retries, :socket_timeout, :total_timeout, :iteration, :deadline
+
+    def initialize(policy, nodes, partition_filter = nil)
+      if partition_filter.nil?
+        return init_for_node(policy, nodes[0]) if nodes.length == 1
+        return init_for_nodes(policy, nodes)
+      end
+
+      # Validate here instead of initial PartitionFilter constructor because total number of
+      # cluster partitions may change on the server and PartitionFilter will never have access
+      # to Cluster instance.  Use fixed number of partitions for now.
+      unless partition_filter.partition_begin.between?(0, Node::PARTITIONS - 1)
+        raise Aerospike::Exceptions::Aerospike.new(
+          Aerospike::ResultCode::PARAMETER_ERROR,
+          "Invalid partition begin #{partition_filter.partition_begin}. Valid range: 0-#{Aerospike::Node::PARTITIONS - 1}"
+        )
+      end
+
+      if partition_filter.count <= 0
+        raise Aerospike::Exceptions::Aerospike.new(
+          Aerospike::ResultCode::PARAMETER_ERROR,
+          "Invalid partition count #{partition_filter.count}"
+        )
+      end
+
+      if partition_filter.partition_begin + partition_filter.count > Node::PARTITIONS
+        raise Aerospike::Exceptions::Aerospike.new(
+          Aerospike::ResultCode::PARAMETER_ERROR,
+          "Invalid partition range (#{partition_filter.partition_begin}, #{partition_filter.partition_begin + partition_filter.count}"
+        )
+      end
+
+      @partition_begin = partition_filter.partition_begin
+      @node_capacity = nodes.length
+      @node_filter = nil
+      @partitions_capacity = partition_filter.count
+      @max_records = policy.max_records
+      @iteration = 1
+
+      if partition_filter.partitions.nil?  then
+        partition_filter.partitions = init_partitions(policy, partition_filter.count, partition_filter.digest)
+      elsif policy.max_records <= 0
+        # Retry all partitions when max_records not specified.
+        partition_filter.partitions.each do |ps|
+          ps.retry = true
+        end
+      end
+
+      @partitions = partition_filter.partitions
+      @partition_filter = partition_filter
+      init_timeout(policy)
+    end
+
+    def assign_partitions_to_nodes(cluster, namespace)
+      list = []
+
+      pmap = cluster.partitions
+      replica_array = pmap[namespace]
+      raise Aerospike::Exceptions::InvalidNamespace("namespace not found in the partition map") if !replica_array
+
+      master = (replica_array.get)[0]
+      master = master.get
+
+      @partitions.each do |part|
+        if part&.retry
+          node = master[part.id]
+
+          unless node
+            raise Exceptions::Aerospike.new(Aerospike::ResultCode::INVALID_NAMESPACE, "Invalid Partition Id #{part.id} for namespace `#{namespace}` in Partition Scan")
+          end
+
+        part.retry = false
+
+
+        # Use node name to check for single node equality because
+        # partition map may be in transitional state between
+        # the old and new node with the same name.
+        next if @node_filter && @node_filter.name != node.name
+
+        np = find_node(list, node)
+
+        unless np
+          # If the partition map is in a transitional state, multiple
+          # node_partitions instances (each with different partitions)
+          # may be created for a single node.
+          np = NodePartitions.new(node)
+          list << np
+        end
+        np.add_partition(part)
+      end
+    end
+
+    if @max_records.positive?
+      # Distribute max_records across nodes.
+      node_size = list.length
+
+      if @max_records < node_size
+        # Only include nodes that have at least 1 record requested.
+        node_size = @max_records
+        list = list[0...node_size]
+      end
+
+      max = 0
+      max = @max_records / node_size if node_size.positive?
+      rem = @max_records - (max * node_size)
+
+      list[0...node_size].each_with_index do |np, i|
+        np.record_max = (i < rem ? max + 1 : max)
+      end
+    end
+
+      @node_partitions_list = list
+      list
+    end
+
+    def init_timeout(policy)
+      @sleep_between_retries = policy.sleep_between_retries
+      @socket_timeout = policy.socket_timeout
+      @total_timeout = policy.timeout
+      if @total_timeout.positive?
+        @deadline = Time.now + @total_timeout
+        if !@socket_timeout || @socket_timeout > @total_timeout
+          @socket_timeout = @total_timeout
+        end
+      end
+    end
+
+    def init_partitions(policy, partition_count, digest)
+      parts_all = Array.new(partition_count)
+
+      (0...partition_count).each do |i|
+        parts_all[i] = Aerospike::PartitionStatus.new(@partition_begin + i)
+      end
+
+      parts_all[0].digest = digest if digest
+
+      @sleep_between_retries = policy.sleep_between_retries
+      @socket_timeout = policy.socket_timeout
+      @total_timeout = policy.timeout
+
+      if @total_timeout.positive?
+        @deadline = Time.now + @total_timeout
+
+        if @socket_timeout == 0 || @socket_timeout > @total_timeout
+          @socket_timeout = @total_timeout
+        end
+      end
+
+      parts_all
+    end
+
+    attr_writer :sleep_between_retries
+
+
+    def find_node(list, node)
+      list.each do |node_partition|
+          # Use pointer equality for performance.
+          return node_partition if node_partition.node == node
+      end
+      nil
+    end
+
+    def partition_unavailable(node_partitions, partition_id)
+      @partitions[partition_id-@partition_begin].retry = true
+      node_partitions.parts_unavailable+=1
+    end
+
+    def set_digest(node_partitions, key)
+      partition_id = key.partition_id
+      @partitions[partition_id-@partition_begin].digest = key.digest
+      node_partitions.record_count+=1
+    end
+
+    def set_last(node_partitions, key, bval)
+      partition_id = key.partition_id()
+      if partition_id-@partition_begin < 0
+        raise "key.partition_id: #{@partition_id}, partition_begin: #{@partition_begin}"
+      end
+      ps = @partitions[partition_id-@partition_begin]
+      ps.digest = key.digest
+      ps.bval = bval
+      node_partitions.record_count+=1
+    end
+
+    def complete?(cluster, policy)
+      record_count = 0
+      parts_unavailable = 0
+
+      @node_partitions_list.each do |np|
+        record_count += np.record_count
+        parts_unavailable += np.parts_unavailable
+      end
+
+      if parts_unavailable == 0
+        if @max_records <= 0
+          @partition_filter&.done = true
+        else
+          if cluster.supports_partition_query.get()
+           done = true
+
+           @node_partitions_list.each do |np|
+             if np.record_count >= np.record_max
+               mark_retry(np)
+               done = false
+             end
+           end
+
+            @partition_filter&.done = done
+          else
+            # Server version >= 6.0 will return all records for each node up to
+            # that node's max. If node's record count reached max, there stilthen
+            # may be records available for that node.
+            @node_partitions_list.each do |np|
+              mark_retry(np) if np.record_count > 0
+            end
+            # Servers version < 6.0 can return less records than max and still
+            # have more records for each node, so the node is only done if nthen
+            # records were retrieved for that node.
+
+            @partition_filter&.done = (record_count == 0)
+          end
+        end
+        return true
+      end
+
+      return true if @max_records&.positive? && record_count >= @max_records
+
+      # Check if limits have been reached
+      if policy.max_retries.positive? && @iteration > policy.max_retries
+        raise Aerospike::Exceptions::Aerospike.new(Aerospike::ResultCode::MAX_RETRIES_EXCEEDED, "Max retries exceeded: #{policy.max_retries}")
+      end
+
+      if policy.total_timeout > 0
+        # Check for total timeout.
+        remaining = @deadline - Time.now - @sleep_between_retries
+
+        raise Aerospike::Exceptions::Timeout.new(policy.totle_timeout, @iteration) if remaining <= 0
+
+        if remaining < @total_timeout
+          @total_timeout = remaining
+
+          if @socket_timeout > @total_timeout
+            @socket_timeout = @total_timeout
+          end
+        end
+      end
+
+      # Prepare for next iteration.
+      if @max_records > 0
+        @max_records -= record_count
+      end
+      @iteration+=1
+      false
+    end
+
+    def should_retry(node_partitions, err)
+      case err
+      when Aerospike::Exceptions::Aerospike
+        case err.result_code
+        when Aerospike::ResultCode::TIMEOUT,
+          Aerospike::ResultCode::NETWORK_ERROR,
+          Aerospike::ResultCode::SERVER_NOT_AVAILABLE,
+          Aerospike::ResultCode::INDEX_NOTFOUND
+            mark_retry(node_partitions)
+            node_partitions.parts_unavailable = node_partitions.parts_full.length + node_partitions.parts_partial.length
+            true
+        end
+      else
+        false
+      end
+    end
+
+    def mark_retry(node_partitions)
+      node_partitions.parts_full.each do |ps|
+        ps.retry = true
+      end
+
+      node_partitions.parts_partial.each do |ps|
+        ps.retry = true
+      end
+    end
+
+    def to_s
+      sb = StringIO.new
+      @partitions.each_with_index do |ps, i|
+        sb << ps.to_s
+        sb << if (i+1)%16 == 0
+          "\n"
+              else
+          "\t"
+              end
+      end
+      sb.string
+    end
+
+    private
+
+    def init_for_nodes(policy, nodes)
+      ppn = Aerospike::Node::PARTITIONS / nodes.length
+      ppn += ppn / 4
+
+      @partition_begin = 0
+      @node_capacity = nodes.length
+      @node_filter = nil
+      @partitions_capacity = ppn
+      @max_records = policy.max_records
+      @iteration = 1
+
+      @partitions = init_partitions(policy, Aerospike::Node::PARTITIONS, nil)
+      init_timeout(policy)
+    end
+
+    def init_for_node(policy, node)
+      @partition_begin = 0
+      @node_capacity = 1
+      @node_filter = node
+      @partitions_capacity = Aerospike::Node::PARTITIONS
+      @max_records = policy.max_records
+      @iteration = 1
+
+      @partitions = init_partitions(policy, Aerospike::Node::PARTITIONS, nil)
+      init_timeout(policy)
+    end
+
+  end
+end

data/lib/aerospike/query/query_command.rb

@@ -82,7 +82,7 @@ module Aerospike
           @data_offset += binNameSize
           fieldCount+=1
         end
-      else    
+      else
         @data_offset += @partitions.length * 2 + FIELD_HEADER_SIZE
         fieldCount += 1
 

data/lib/aerospike/query/query_executor.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+# Copyright 2014-2020 Aerospike, Inc.
+#
+# Portions may be licensed to Aerospike, Inc. under one or more contributor
+# license agreements.
+#
+# Licensed 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.
+
+module Aerospike
+  class QueryExecutor # :nodoc:
+
+    def self.query_partitions(cluster, policy, tracker, statement, recordset)
+      interval = policy.sleep_between_retries
+
+      should_retry = false
+
+      loop do
+        list = tracker.assign_partitions_to_nodes(cluster, statement.namespace)
+
+        if policy.concurrent_nodes
+          threads = []
+          # Use a thread per node
+          list.each do |node_partition|
+
+            threads << Thread.new do
+              Thread.current.abort_on_exception = true
+              command = QueryPartitionCommand.new(node_partition.node, tracker, policy, statement, recordset, node_partition)
+              begin
+                command.execute
+              rescue => e
+                should_retry ||= command.should_retry(e)
+                # puts "should retry: #{should_retry}"
+                Aerospike.logger.error(e.backtrace.join("\n")) unless e == QUERY_TERMINATED_EXCEPTION
+              end
+            end
+          end
+          threads.each(&:join)
+        else
+          # Use a single thread for all nodes for all node
+          list.each do |node_partition|
+            command = QueryPartitionCommand.new(node_partition.node, tracker, policy, statement, recordset, node_partition)
+            begin
+              command.execute
+            rescue => e
+              should_retry ||= command.should_retry(e)
+              Aerospike.logger.error(e.backtrace.join("\n")) unless e == QUERY_TERMINATED_EXCEPTION
+            end
+          end
+        end
+
+        complete = tracker.complete?(@cluster, policy)
+
+        if complete || !should_retry
+          recordset.thread_finished
+          return
+        end
+        sleep(interval) if policy.sleep_between_retries > 0
+        statement.reset_task_id
+      end
+    end
+
+  end
+
+end

data/lib/aerospike/query/query_partition_command.rb

@@ -0,0 +1,266 @@
+# encoding: utf-8
+# Copyright 2014-2020 Aerospike, Inc.
+#
+# Portions may be licensed to Aerospike, Inc. under one or more contributor
+# license agreements.
+#
+# Licensed 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.
+
+require 'aerospike/query/stream_command'
+require 'aerospike/query/recordset'
+
+module Aerospike
+
+  private
+
+  class QueryPartitionCommand < QueryCommand #:nodoc:
+
+    def initialize(node, tracker, policy, statement, recordset, node_partitions)
+      super(node, policy, statement, recordset, @node_partitions)
+      @node_partitions = node_partitions
+      @tracker = tracker
+    end
+
+    def write_buffer
+      function_arg_buffer = nil
+      field_count = 0
+      filter_size = 0
+      bin_name_size = 0
+
+      begin_cmd
+
+      if @statement.namespace
+        @data_offset += @statement.namespace.bytesize + FIELD_HEADER_SIZE
+        field_count+=1
+      end
+
+      if @statement.set_name
+        @data_offset += @statement.set_name.bytesize + FIELD_HEADER_SIZE
+        field_count+=1
+      end
+
+      # Estimate recordsPerSecond field size. This field is used in new servers and not used
+      # (but harmless to add) in old servers.
+      if @policy.records_per_second > 0
+        @data_offset += 4 + FIELD_HEADER_SIZE
+        field_count+=1
+      end
+
+      # Estimate socket timeout field size. This field is used in new servers and not used
+      # (but harmless to add) in old servers.
+      @data_offset += 4 + FIELD_HEADER_SIZE
+      field_count+=1
+
+      # Estimate task_id field.
+      @data_offset += 8 + FIELD_HEADER_SIZE
+      field_count+=1
+
+      filter = @statement.filters[0]
+      bin_names = @statement.bin_names
+      packed_ctx = nil
+
+      if filter
+        col_type = filter.collection_type
+
+        # Estimate INDEX_TYPE field.
+        if col_type > 0
+          @data_offset += FIELD_HEADER_SIZE + 1
+          field_count+=1
+        end
+
+        # Estimate INDEX_RANGE field.
+        @data_offset += FIELD_HEADER_SIZE
+        filter_size+=1  # num filters
+        filter_size += filter.estimate_size
+
+        @data_offset += filter_size
+        field_count+=1
+
+        # TODO: Implement
+        # packed_ctx = filter.packed_ctx
+        # if packed_ctx
+        #   @data_offset += FIELD_HEADER_SIZE + packed_ctx.length
+        #   field_count+=1
+        # end
+      end
+
+      @statement.set_task_id
+      predexp = @policy.predexp || @statement.predexp
+
+      if predexp
+        @data_offset += FIELD_HEADER_SIZE
+        pred_size = Aerospike::PredExp.estimate_size(predexp)
+        @data_offset += pred_size
+        field_count += 1
+      end
+
+      # Estimate aggregation/background function size.
+      if @statement.function_name
+        @data_offset += FIELD_HEADER_SIZE + 1 # udf type
+        @data_offset += @statement.package_name.bytesize + FIELD_HEADER_SIZE
+        @data_offset += @statement.function_name.bytesize + FIELD_HEADER_SIZE
+
+        function_arg_buffer=''
+        if @statement.function_args && @statement.function_args.length > 0
+          function_arg_buffer = Value.of(@statement.function_args).to_bytes
+        end
+        @data_offset += FIELD_HEADER_SIZE + function_arg_buffer.bytesize
+        field_count += 4
+      end
+
+      max_records = 0
+      parts_full_size = 0
+      parts_partial_digest_size = 0
+      parts_partial_bval_size = 0
+
+      unless @node_partitions.nil?
+        parts_full_size = @node_partitions.parts_full.length * 2
+        parts_partial_digest_size = @node_partitions.parts_partial.length * 20
+
+        unless filter.nil?
+          parts_partial_bval_size = @node_partitions.parts_partial.length * 8
+        end
+        max_records = @node_partitions.record_max
+      end
+
+      if parts_full_size > 0
+        @data_offset += parts_full_size + FIELD_HEADER_SIZE
+        field_count+=1
+      end
+
+      if parts_partial_digest_size > 0
+        @data_offset += parts_partial_digest_size + FIELD_HEADER_SIZE
+        field_count+=1
+      end
+
+      if parts_partial_bval_size > 0
+        @data_offset += parts_partial_bval_size + FIELD_HEADER_SIZE
+        field_count+=1
+      end
+
+      # Estimate max records field size. This field is used in new servers and not used
+      # (but harmless to add) in old servers.
+      if max_records > 0
+        @data_offset += 8 + FIELD_HEADER_SIZE
+        field_count+=1
+      end
+
+      operation_count = 0
+      unless bin_names.empty?
+        # Estimate size for selected bin names (query bin names already handled for old servers).
+        bin_names.each do |bin_name|
+          estimate_operation_size_for_bin_name(bin_name)
+        end
+        operation_count = bin_names.length
+      end
+
+      projected_offset = @data_offset
+
+      size_buffer
+
+      read_attr = INFO1_READ
+      read_attr |= INFO1_NOBINDATA if !@policy.include_bin_data
+      read_attr |= INFO1_SHORT_QUERY if @policy.short_query
+
+      infoAttr = INFO3_PARTITION_DONE
+
+      write_header(@policy, read_attr, 0, field_count, operation_count)
+
+      write_field_string(@statement.namespace, FieldType::NAMESPACE) if @statement.namespace
+      write_field_string(@statement.set_name, FieldType::TABLE) if @statement.set_name
+
+      # Write records per second.
+      write_field_int(@policy.records_per_second, FieldType::RECORDS_PER_SECOND) if @policy.records_per_second > 0
+
+      # Write socket idle timeout.
+      write_field_int(@policy.socket_timeout, FieldType::SOCKET_TIMEOUT)
+
+      # Write task_id field
+      write_field_int64(@statement.task_id, FieldType::TRAN_ID)
+
+      unless predexp.nil?
+        write_field_header(pred_size, Aerospike::FieldType::PREDEXP)
+        @data_offset = Aerospike::PredExp.write(
+          predexp, @data_buffer, @data_offset
+        )
+      end
+
+      if filter
+        type = filter.collection_type
+
+        if type > 0
+          write_field_header(1, FieldType::INDEX_TYPE)
+          @data_offset += @data_buffer.write_byte(type, @data_offset)
+        end
+
+        write_field_header(filter_size, FieldType::INDEX_RANGE)
+        @data_offset += @data_buffer.write_byte(1, @data_offset)
+        @data_offset = filter.write(@data_buffer, @data_offset)
+
+        # TODO: Implement
+        # if packed_ctx
+        #   write_field_header(packed_ctx.length, FieldType::INDEX_CONTEXT)
+        #   @data_buffer.write_binary(packed_ctx, @data_offset)
+        # end
+      end
+
+      if @statement.function_name
+        write_field_header(1, FieldType::UDF_OP)
+        @data_offset += @data_buffer.write_byte(1, @data_offset)
+        write_field_string(@statement.package_name, FieldType::UDF_PACKAGE_NAME)
+        write_field_string(@statement.function_name, FieldType::UDF_FUNCTION)
+        write_field_string(function_arg_buffer, FieldType::UDF_ARGLIST)
+      end
+
+      if parts_full_size > 0
+        write_field_header(parts_full_size, FieldType::PID_ARRAY)
+        @node_partitions.parts_full.each do |part|
+          @data_offset += @data_buffer.write_uint16_little_endian(part.id, @data_offset)
+        end
+      end
+
+      if parts_partial_digest_size > 0
+        write_field_header(parts_partial_digest_size, FieldType::DIGEST_ARRAY)
+        @node_partitions.parts_partial.each do |part|
+          @data_offset += @data_buffer.write_binary(part.digest, @data_offset)
+        end
+      end
+
+      if parts_partial_bval_size > 0
+        write_field_header(parts_partial_bval_size, FieldType::BVAL_ARRAY)
+        @node_partitions.parts_partial.each do |part|
+          @data_offset += @data_buffer.write_uint64_little_endian(part.bval, @data_offset)
+        end
+      end
+
+      if max_records > 0
+        write_field(max_records, FieldType::MAX_RECORDS)
+      end
+
+      unless bin_names.empty?
+        bin_names.each do |bin_name|
+          write_operation_for_bin_name(bin_name, Operation::READ)
+        end
+      end
+
+      end_cmd
+
+      nil
+    end
+
+    def should_retry(e)
+      # !! converts nil to false
+      !!@tracker&.should_retry(@node_partitions, e)
+    end
+
+  end # class
+
+end # module