From 02245ce41f5c47f53c6a5e4db6560c6e5381e853 Mon Sep 17 00:00:00 2001
From: Daniel Waterworth <me@danielwaterworth.com>
Date: Fri, 10 Dec 2021 14:25:26 -0600
Subject: [PATCH] PERF: Redis snapshotting during tests (#15260)

We can fake redis transactions so that `fab!` works for redis and PG
data, but it's too slow to be used indiscriminately. Instead, you can
opt into it with the `use_redis_snapshotting` helper.

Insofar as snapshotting allows us to `fab!` more things, it provides a
speedup.
---
 lib/discourse_redis.rb                  |  3 +-
 lib/redis_snapshot.rb                   | 41 +++++++++++++++++++++++++
 spec/helpers/redis_snapshot_helper.rb   | 21 +++++++++++++
 spec/rails_helper.rb                    |  2 ++
 spec/requests/topics_controller_spec.rb | 10 +++---
 5 files changed, 72 insertions(+), 5 deletions(-)
 create mode 100644 lib/redis_snapshot.rb
 create mode 100644 spec/helpers/redis_snapshot_helper.rb

diff --git a/lib/discourse_redis.rb b/lib/discourse_redis.rb
index 1ed8fc6a5ed..869ac8ace47 100644
--- a/lib/discourse_redis.rb
+++ b/lib/discourse_redis.rb
@@ -53,7 +53,8 @@ class DiscourseRedis
    :msetnx, :persist, :pexpire, :pexpireat, :psetex, :pttl, :rename, :renamenx, :rpop, :rpoplpush, :rpush, :rpushx, :sadd, :scard,
    :sdiff, :set, :setbit, :setex, :setnx, :setrange, :sinter, :sismember, :smembers, :sort, :spop, :srandmember, :srem, :strlen,
    :sunion, :ttl, :type, :watch, :zadd, :zcard, :zcount, :zincrby, :zrange, :zrangebyscore, :zrank, :zrem, :zremrangebyrank,
-   :zremrangebyscore, :zrevrange, :zrevrangebyscore, :zrevrank, :zrangebyscore ].each do |m|
+   :zremrangebyscore, :zrevrange, :zrevrangebyscore, :zrevrank, :zrangebyscore,
+   :dump, :restore].each do |m|
     define_method m do |*args, **kwargs|
       args[0] = "#{namespace}:#{args[0]}" if @namespace
       DiscourseRedis.ignore_readonly { @redis.public_send(m, *args, **kwargs) }
diff --git a/lib/redis_snapshot.rb b/lib/redis_snapshot.rb
new file mode 100644
index 00000000000..3eb78fb3a16
--- /dev/null
+++ b/lib/redis_snapshot.rb
@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+class RedisSnapshot
+  def self.begin_faux_transaction(redis = Discourse.redis)
+    @stack ||= []
+    @stack.push(RedisSnapshot.load(redis))
+  end
+
+  def self.end_faux_transaction(redis = Discourse.redis)
+    @stack.pop.restore(redis)
+  end
+
+  def self.load(redis = Discourse.redis)
+    keys = redis.keys
+
+    values =
+      redis.pipelined do
+        keys.each do |key|
+          redis.dump(key)
+        end
+      end
+
+    new(keys.zip(values).delete_if { |k, v| v.nil? })
+  end
+
+  def initialize(dump)
+    @dump = dump
+  end
+
+  def restore(redis = Discourse.redis)
+    redis.pipelined do
+      redis.flushdb
+
+      @dump.each do |key, value|
+        redis.restore(key, 0, value)
+      end
+    end
+
+    nil
+  end
+end
diff --git a/spec/helpers/redis_snapshot_helper.rb b/spec/helpers/redis_snapshot_helper.rb
new file mode 100644
index 00000000000..fd8f531b083
--- /dev/null
+++ b/spec/helpers/redis_snapshot_helper.rb
@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module RedisSnapshotHelper
+  def use_redis_snapshotting
+    before(:all) do
+      RedisSnapshot.begin_faux_transaction
+    end
+
+    after(:all) do
+      RedisSnapshot.end_faux_transaction
+    end
+
+    before(:each) do
+      RedisSnapshot.begin_faux_transaction
+    end
+
+    after(:each) do
+      RedisSnapshot.end_faux_transaction
+    end
+  end
+end
diff --git a/spec/rails_helper.rb b/spec/rails_helper.rb
index a5deadd6302..7f96bd11521 100644
--- a/spec/rails_helper.rb
+++ b/spec/rails_helper.rb
@@ -78,6 +78,7 @@ end
 # in spec/support/ and its subdirectories.
 Dir[Rails.root.join("spec/support/**/*.rb")].each { |f| require f }
 Dir[Rails.root.join("spec/fabricators/*.rb")].each { |f| require f }
+require_relative './helpers/redis_snapshot_helper'
 
 # Require plugin helpers at plugin/[plugin]/spec/plugin_helper.rb (includes symlinked plugins).
 if ENV['LOAD_PLUGINS'] == "1"
@@ -181,6 +182,7 @@ end
 RSpec.configure do |config|
   config.fail_fast = ENV['RSPEC_FAIL_FAST'] == "1"
   config.silence_filter_announcements = ENV['RSPEC_SILENCE_FILTER_ANNOUNCEMENTS'] == "1"
+  config.extend RedisSnapshotHelper
   config.include Helpers
   config.include MessageBus
   config.include RSpecHtmlMatchers
diff --git a/spec/requests/topics_controller_spec.rb b/spec/requests/topics_controller_spec.rb
index 875fc507338..b53eda3a063 100644
--- a/spec/requests/topics_controller_spec.rb
+++ b/spec/requests/topics_controller_spec.rb
@@ -1756,11 +1756,13 @@ RSpec.describe TopicsController do
   end
 
   describe '#show' do
-    fab!(:private_topic) { Fabricate(:private_message_topic) }
-    let!(:topic) { Fabricate(:post).topic }
+    use_redis_snapshotting
 
-    let!(:p1) { Fabricate(:post, user: topic.user) }
-    let!(:p2) { Fabricate(:post, user: topic.user) }
+    fab!(:private_topic) { Fabricate(:private_message_topic) }
+    fab!(:topic) { Fabricate(:post).topic }
+
+    fab!(:p1) { Fabricate(:post, user: topic.user) }
+    fab!(:p2) { Fabricate(:post, user: topic.user) }
 
     describe 'when topic is not allowed' do
       it 'should return the right response' do