10 files changed, 351 insertions, 3 deletions

diff --git a/lib/api/helpers/pagination.rb b/lib/api/helpers/pagination.rb
index 9c5b355e823..642053949d9 100644
--- a/lib/api/helpers/pagination.rb
+++ b/lib/api/helpers/pagination.rb
@@ -4,7 +4,21 @@ module API
   module Helpers
     module Pagination
       def paginate(relation)
-        ::Gitlab::Pagination::OffsetPagination.new(self).paginate(relation)
+        return Gitlab::Pagination::OffsetPagination.new(self).paginate(relation) unless keyset_pagination_enabled?
+
+        request_context = Gitlab::Pagination::Keyset::RequestContext.new(self)
+
+        unless Gitlab::Pagination::Keyset.available?(request_context, relation)
+          return error!('Keyset pagination is not yet available for this type of request', 501)
+        end
+
+        Gitlab::Pagination::Keyset.paginate(request_context, relation)
+      end
+
+      private
+
+      def keyset_pagination_enabled?
+        params[:pagination] == 'keyset' && Feature.enabled?(:api_keyset_pagination)
       end
     end
   end
diff --git a/lib/api/projects.rb b/lib/api/projects.rb
index a1fce9e8b20..666bd2771f9 100644
--- a/lib/api/projects.rb
+++ b/lib/api/projects.rb
@@ -82,7 +82,6 @@ module API
       def present_projects(projects, options = {})
         projects = reorder_projects(projects)
         projects = apply_filters(projects)
-        projects = paginate(projects)
         projects, options = with_custom_attributes(projects, options)
 
         options = options.reverse_merge(
@@ -93,7 +92,10 @@ module API
         )
         options[:with] = Entities::BasicProjectDetails if params[:simple]
 
-        present options[:with].prepare_relation(projects, options), options
+        projects = options[:with].prepare_relation(projects, options)
+        projects = paginate(projects)
+
+        present projects, options
       end
 
       def translate_params_for_compatibility(params)
diff --git a/lib/gitlab/marginalia.rb b/lib/gitlab/marginalia.rb
new file mode 100644
index 00000000000..d2e0e335127
--- /dev/null
+++ b/lib/gitlab/marginalia.rb
@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Marginalia
+    MARGINALIA_FEATURE_FLAG = :marginalia
+
+    def self.set_application_name
+      ::Marginalia.application_name = Gitlab.process_name
+    end
+
+    def self.enable_sidekiq_instrumentation
+      if Sidekiq.server?
+        ::Marginalia::SidekiqInstrumentation.enable!
+      end
+    end
+
+    def self.feature_enabled?
+      return false unless Gitlab::Database.cached_table_exists?('features')
+
+      Feature.enabled?(MARGINALIA_FEATURE_FLAG)
+    end
+  end
+end
diff --git a/lib/gitlab/marginalia/active_record_instrumentation.rb b/lib/gitlab/marginalia/active_record_instrumentation.rb
new file mode 100644
index 00000000000..f4500a48090
--- /dev/null
+++ b/lib/gitlab/marginalia/active_record_instrumentation.rb
@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+# Patch to annotate sql only when the feature is enabled.
+module Gitlab
+  module Marginalia
+    module ActiveRecordInstrumentation
+      # CAUTION:
+      # Any method call which generates a query inside this function will get into a recursive loop unless called within `Marginalia.without_annotation` method.
+      def annotate_sql(sql)
+        if ActiveRecord::Base.connected? &&
+          ::Marginalia.annotation_allowed? &&
+          ::Marginalia.without_annotation { Gitlab::Marginalia.feature_enabled? }
+          super(sql)
+        else
+          sql
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/marginalia/comment.rb b/lib/gitlab/marginalia/comment.rb
new file mode 100644
index 00000000000..a0eee823763
--- /dev/null
+++ b/lib/gitlab/marginalia/comment.rb
@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+# Module to support correlation_id and additional job details.
+module Gitlab
+  module Marginalia
+    module Comment
+      private
+
+      def jid
+        bg_job["jid"] if bg_job.present?
+      end
+
+      def job_class
+        bg_job["class"] if bg_job.present?
+      end
+
+      def correlation_id
+        if bg_job.present?
+          bg_job["correlation_id"]
+        else
+          Labkit::Correlation::CorrelationId.current_id
+        end
+      end
+
+      def bg_job
+        job = ::Marginalia::Comment.marginalia_job
+
+        # We are using 'Marginalia::SidekiqInstrumentation' which does not support 'ActiveJob::Base'.
+        # Gitlab also uses 'ActionMailer::DeliveryJob' which inherits from ActiveJob::Base.
+        # So below condition is used to return metadata for such jobs.
+        if job && job.is_a?(ActionMailer::DeliveryJob)
+          {
+            "class" => job.arguments.first,
+            "jid"   => job.job_id
+          }
+        else
+          job
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/marginalia/inline_annotation.rb b/lib/gitlab/marginalia/inline_annotation.rb
new file mode 100644
index 00000000000..6d78f3b81ec
--- /dev/null
+++ b/lib/gitlab/marginalia/inline_annotation.rb
@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+# Module with util methods to support ::Marginalia.without_annotation method.
+
+module Gitlab
+  module Marginalia
+    module InlineAnnotation
+      def without_annotation(&block)
+        return unless block.present?
+
+        annotation_stack.push(false)
+        yield
+      ensure
+        annotation_stack.pop
+      end
+
+      def with_annotation(comment, &block)
+        annotation_stack.push(true)
+        super(comment, &block)
+      ensure
+        annotation_stack.pop
+      end
+
+      def annotation_stack
+        Thread.current[:annotation_stack] ||= []
+      end
+
+      def annotation_stack_top
+        annotation_stack.last
+      end
+
+      def annotation_allowed?
+        annotation_stack.empty? ? true : annotation_stack_top
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/pagination/keyset.rb b/lib/gitlab/pagination/keyset.rb
new file mode 100644
index 00000000000..5bd45fa9b56
--- /dev/null
+++ b/lib/gitlab/pagination/keyset.rb
@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Pagination
+    module Keyset
+      def self.paginate(request_context, relation)
+        Gitlab::Pagination::Keyset::Pager.new(request_context).paginate(relation)
+      end
+
+      def self.available?(request_context, relation)
+        order_by = request_context.page.order_by
+
+        # This is only available for Project and order-by id (asc/desc)
+        return false unless relation.klass == Project
+        return false unless order_by.size == 1 && order_by[:id]
+
+        true
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/pagination/keyset/page.rb b/lib/gitlab/pagination/keyset/page.rb
new file mode 100644
index 00000000000..3f71822a7c7
--- /dev/null
+++ b/lib/gitlab/pagination/keyset/page.rb
@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Pagination
+    module Keyset
+      # A Page models the pagination information for a particular page of the collection
+      class Page
+        # Default and maximum size of records for a page
+        DEFAULT_PAGE_SIZE = 20
+
+        attr_accessor :lower_bounds, :end_reached
+        attr_reader :order_by
+
+        def initialize(order_by: {}, lower_bounds: nil, per_page: DEFAULT_PAGE_SIZE, end_reached: false)
+          @order_by = order_by.symbolize_keys
+          @lower_bounds = lower_bounds&.symbolize_keys
+          @per_page = per_page
+          @end_reached = end_reached
+        end
+
+        # Number of records to return per page
+        def per_page
+          return DEFAULT_PAGE_SIZE if @per_page <= 0
+
+          [@per_page, DEFAULT_PAGE_SIZE].min
+        end
+
+        # Determine whether this page indicates the end of the collection
+        def end_reached?
+          @end_reached
+        end
+
+        # Construct a Page for the next page
+        # Uses identical order_by/per_page information for the next page
+        def next(lower_bounds, end_reached)
+          dup.tap do |next_page|
+            next_page.lower_bounds = lower_bounds&.symbolize_keys
+            next_page.end_reached = end_reached
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/pagination/keyset/pager.rb b/lib/gitlab/pagination/keyset/pager.rb
new file mode 100644
index 00000000000..99b125cc2a0
--- /dev/null
+++ b/lib/gitlab/pagination/keyset/pager.rb
@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Pagination
+    module Keyset
+      class Pager
+        attr_reader :request
+
+        def initialize(request)
+          @request = request
+        end
+
+        def paginate(relation)
+          # Validate assumption: The last two columns must match the page order_by
+          validate_order!(relation)
+
+          # This performs the database query and retrieves records
+          # We retrieve one record more to check if we have data beyond this page
+          all_records = relation.limit(page.per_page + 1).to_a # rubocop: disable CodeReuse/ActiveRecord
+
+          records_for_page = all_records.first(page.per_page)
+
+          # If we retrieved more records than belong on this page,
+          # we know there's a next page
+          there_is_more = all_records.size > records_for_page.size
+          apply_headers(records_for_page.last, there_is_more)
+
+          records_for_page
+        end
+
+        private
+
+        def apply_headers(last_record_in_page, there_is_more)
+          end_reached = last_record_in_page.nil? || !there_is_more
+          lower_bounds = last_record_in_page&.slice(page.order_by.keys)
+
+          next_page = page.next(lower_bounds, end_reached)
+
+          request.apply_headers(next_page)
+        end
+
+        def page
+          @page ||= request.page
+        end
+
+        def validate_order!(rel)
+          present_order = rel.order_values.map { |val| [val.expr.name.to_sym, val.direction] }.last(2).to_h
+
+          unless page.order_by == present_order
+            raise ArgumentError, "Page's order_by does not match the relation's order: #{present_order} vs #{page.order_by}"
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/pagination/keyset/request_context.rb b/lib/gitlab/pagination/keyset/request_context.rb
new file mode 100644
index 00000000000..aeaed7587b3
--- /dev/null
+++ b/lib/gitlab/pagination/keyset/request_context.rb
@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Pagination
+    module Keyset
+      class RequestContext
+        attr_reader :request
+
+        DEFAULT_SORT_DIRECTION = :desc
+        PRIMARY_KEY = :id
+
+        # A tie breaker is added as an additional order-by column
+        # to establish a well-defined order. We use the primary key
+        # column here.
+        TIE_BREAKER = { PRIMARY_KEY => DEFAULT_SORT_DIRECTION }.freeze
+
+        def initialize(request)
+          @request = request
+        end
+
+        # extracts Paging information from request parameters
+        def page
+          @page ||= Page.new(order_by: order_by, per_page: params[:per_page])
+        end
+
+        def apply_headers(next_page)
+          request.header('Links', pagination_links(next_page))
+        end
+
+        private
+
+        def order_by
+          return TIE_BREAKER.dup unless params[:order_by]
+
+          order_by = { params[:order_by].to_sym => params[:sort]&.to_sym || DEFAULT_SORT_DIRECTION }
+
+          # Order by an additional unique key, we use the primary key here
+          order_by = order_by.merge(TIE_BREAKER) unless order_by[PRIMARY_KEY]
+
+          order_by
+        end
+
+        def params
+          @params ||= request.params
+        end
+
+        def lower_bounds_params(page)
+          page.lower_bounds.each_with_object({}) do |(column, value), params|
+            filter = filter_with_comparator(page, column)
+            params[filter] = value
+          end
+        end
+
+        def filter_with_comparator(page, column)
+          direction = page.order_by[column]
+
+          if direction&.to_sym == :desc
+            "#{column}_before"
+          else
+            "#{column}_after"
+          end
+        end
+
+        def page_href(page)
+          base_request_uri.tap do |uri|
+            uri.query = query_params_for(page).to_query
+          end.to_s
+        end
+
+        def pagination_links(next_page)
+          return if next_page.end_reached?
+
+          %(<#{page_href(next_page)}>; rel="next")
+        end
+
+        def base_request_uri
+          @base_request_uri ||= URI.parse(request.request.url).tap do |uri|
+            uri.host = Gitlab.config.gitlab.host
+            uri.port = Gitlab.config.gitlab.port
+          end
+        end
+
+        def query_params_for(page)
+          request.params.merge(lower_bounds_params(page))
+        end
+      end
+    end
+  end
+end