Add middleware for ETag caching with Redis

5 files changed, 269 insertions, 0 deletions

diff --git a/changelogs/unreleased/etag-notes-polling.yml b/changelogs/unreleased/etag-notes-polling.yml
new file mode 100644
index 00000000000..53990821d25
--- /dev/null
+++ b/changelogs/unreleased/etag-notes-polling.yml
@@ -0,0 +1,4 @@
+---
+title: Use ETag to improve performance of issue notes polling
+merge_request: 9036
+author:
diff --git a/config/initializers/etag_caching.rb b/config/initializers/etag_caching.rb
new file mode 100644
index 00000000000..eba88801141
--- /dev/null
+++ b/config/initializers/etag_caching.rb
@@ -0,0 +1,4 @@
+# This middleware has to come after Gitlab::Metrics::RackMiddleware
+# in the middleware stack, because it tracks events with
+# GitLab Performance Monitoring
+Rails.application.config.middleware.use(Gitlab::EtagCaching::Middleware)
diff --git a/lib/gitlab/etag_caching/middleware.rb b/lib/gitlab/etag_caching/middleware.rb
new file mode 100644
index 00000000000..0f24f9bbfde
--- /dev/null
+++ b/lib/gitlab/etag_caching/middleware.rb
@@ -0,0 +1,66 @@
+module Gitlab
+  module EtagCaching
+    class Middleware
+      RESERVED_WORDS = ProjectPathValidator::RESERVED.map { |word| "/#{word}/" }.join('|')
+      ROUTE_REGEXP = Regexp.union(
+        %r(^(?!.*(#{RESERVED_WORDS})).*/noteable/issue/\d+/notes\z)
+      )
+
+      def initialize(app)
+        @app = app
+      end
+
+      def call(env)
+        return @app.call(env) unless enabled_for_current_route?(env)
+        Gitlab::Metrics.add_event(:etag_caching_middleware_used)
+
+        etag, cached_value_present = get_etag(env)
+        if_none_match = env['HTTP_IF_NONE_MATCH']
+
+        if if_none_match == etag
+          Gitlab::Metrics.add_event(:etag_caching_cache_hit)
+          [304, { 'ETag' => etag }, ['']]
+        else
+          track_cache_miss(if_none_match, cached_value_present)
+
+          status, headers, body = @app.call(env)
+          headers['ETag'] = etag
+          [status, headers, body]
+        end
+      end
+
+      private
+
+      def enabled_for_current_route?(env)
+        ROUTE_REGEXP.match(env['PATH_INFO'])
+      end
+
+      def get_etag(env)
+        cache_key = env['PATH_INFO']
+        store = Store.new
+        current_value = store.get(cache_key)
+        cached_value_present = current_value.present?
+
+        unless cached_value_present
+          current_value = store.touch(cache_key, only_if_missing: true)
+        end
+
+        [weak_etag_format(current_value), cached_value_present]
+      end
+
+      def weak_etag_format(value)
+        %Q{W/"#{value}"}
+      end
+
+      def track_cache_miss(if_none_match, cached_value_present)
+        if if_none_match.blank?
+          Gitlab::Metrics.add_event(:etag_caching_header_missing)
+        elsif !cached_value_present
+          Gitlab::Metrics.add_event(:etag_caching_key_not_found)
+        else
+          Gitlab::Metrics.add_event(:etag_caching_resource_changed)
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/etag_caching/store.rb b/lib/gitlab/etag_caching/store.rb
new file mode 100644
index 00000000000..9532e432f78
--- /dev/null
+++ b/lib/gitlab/etag_caching/store.rb
@@ -0,0 +1,32 @@
+module Gitlab
+  module EtagCaching
+    class Store
+      EXPIRY_TIME = 10.minutes
+      REDIS_NAMESPACE = 'etag:'.freeze
+
+      def get(key)
+        Gitlab::Redis.with { |redis| redis.get(redis_key(key)) }
+      end
+
+      def touch(key, only_if_missing: false)
+        etag = generate_etag
+
+        Gitlab::Redis.with do |redis|
+          redis.set(redis_key(key), etag, ex: EXPIRY_TIME, nx: only_if_missing)
+        end
+
+        etag
+      end
+
+      private
+
+      def generate_etag
+        SecureRandom.hex
+      end
+
+      def redis_key(key)
+        "#{REDIS_NAMESPACE}#{key}"
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/etag_caching/middleware_spec.rb b/spec/lib/gitlab/etag_caching/middleware_spec.rb
new file mode 100644
index 00000000000..8b5bfc4dbb0
--- /dev/null
+++ b/spec/lib/gitlab/etag_caching/middleware_spec.rb
@@ -0,0 +1,163 @@
+require 'spec_helper'
+
+describe Gitlab::EtagCaching::Middleware do
+  let(:app) { double(:app) }
+  let(:middleware) { described_class.new(app) }
+  let(:app_status_code) { 200 }
+  let(:if_none_match) { nil }
+  let(:enabled_path) { '/gitlab-org/gitlab-ce/noteable/issue/1/notes' }
+
+  context 'when ETag caching is not enabled for current route' do
+    let(:path) { '/gitlab-org/gitlab-ce/tree/master/noteable/issue/1/notes' }
+
+    before do
+      mock_app_response
+    end
+
+    it 'does not add ETag header' do
+      _, headers, _ = middleware.call(build_env(path, if_none_match))
+
+      expect(headers['ETag']).to be_nil
+    end
+
+    it 'passes status code from app' do
+      status, _, _ = middleware.call(build_env(path, if_none_match))
+
+      expect(status).to eq app_status_code
+    end
+  end
+
+  context 'when there is no ETag in store for given resource' do
+    let(:path) { enabled_path }
+
+    before do
+      mock_app_response
+      mock_value_in_store(nil)
+    end
+
+    it 'generates ETag' do
+      expect_any_instance_of(Gitlab::EtagCaching::Store)
+        .to receive(:touch).and_return('123')
+
+      middleware.call(build_env(path, if_none_match))
+    end
+
+    context 'when If-None-Match header was specified' do
+      let(:if_none_match) { 'W/"abc"' }
+
+      it 'tracks "etag_caching_key_not_found" event' do
+        expect(Gitlab::Metrics).to receive(:add_event)
+          .with(:etag_caching_middleware_used)
+        expect(Gitlab::Metrics).to receive(:add_event)
+          .with(:etag_caching_key_not_found)
+
+        middleware.call(build_env(path, if_none_match))
+      end
+    end
+  end
+
+  context 'when there is ETag in store for given resource' do
+    let(:path) { enabled_path }
+
+    before do
+      mock_app_response
+      mock_value_in_store('123')
+    end
+
+    it 'returns this value as header' do
+      _, headers, _ = middleware.call(build_env(path, if_none_match))
+
+      expect(headers['ETag']).to eq 'W/"123"'
+    end
+  end
+
+  context 'when If-None-Match header matches ETag in store' do
+    let(:path) { enabled_path }
+    let(:if_none_match) { 'W/"123"' }
+
+    before do
+      mock_value_in_store('123')
+    end
+
+    it 'does not call app' do
+      expect(app).not_to receive(:call)
+
+      middleware.call(build_env(path, if_none_match))
+    end
+
+    it 'returns status code 304' do
+      status, _, _ = middleware.call(build_env(path, if_none_match))
+
+      expect(status).to eq 304
+    end
+
+    it 'tracks "etag_caching_cache_hit" event' do
+      expect(Gitlab::Metrics).to receive(:add_event)
+        .with(:etag_caching_middleware_used)
+      expect(Gitlab::Metrics).to receive(:add_event)
+        .with(:etag_caching_cache_hit)
+
+      middleware.call(build_env(path, if_none_match))
+    end
+  end
+
+  context 'when If-None-Match header does not match ETag in store' do
+    let(:path) { enabled_path }
+    let(:if_none_match) { 'W/"abc"' }
+
+    before do
+      mock_value_in_store('123')
+    end
+
+    it 'calls app' do
+      expect(app).to receive(:call).and_return([app_status_code, {}, ['body']])
+
+      middleware.call(build_env(path, if_none_match))
+    end
+
+    it 'tracks "etag_caching_resource_changed" event' do
+      mock_app_response
+
+      expect(Gitlab::Metrics).to receive(:add_event)
+        .with(:etag_caching_middleware_used)
+      expect(Gitlab::Metrics).to receive(:add_event)
+        .with(:etag_caching_resource_changed)
+
+      middleware.call(build_env(path, if_none_match))
+    end
+  end
+
+  context 'when If-None-Match header is not specified' do
+    let(:path) { enabled_path }
+
+    before do
+      mock_value_in_store('123')
+      mock_app_response
+    end
+
+    it 'tracks "etag_caching_header_missing" event' do
+      expect(Gitlab::Metrics).to receive(:add_event)
+        .with(:etag_caching_middleware_used)
+      expect(Gitlab::Metrics).to receive(:add_event)
+        .with(:etag_caching_header_missing)
+
+      middleware.call(build_env(path, if_none_match))
+    end
+  end
+
+  def mock_app_response
+    allow(app).to receive(:call).and_return([app_status_code, {}, ['body']])
+  end
+
+  def mock_value_in_store(value)
+    allow_any_instance_of(Gitlab::EtagCaching::Store)
+      .to receive(:get).and_return(value)
+  end
+
+  def build_env(path, if_none_match)
+    {
+      'PATH_INFO' => path,
+      'HTTP_IF_NONE_MATCH' => if_none_match
+    }
+  end
+end