Add a statement timeout to queries [PIMCORE-485] (#31)

Author: erikkessler1

CHANGELOG.md

@@ -1,5 +1,11 @@
 # postgres-vacuum-monitor
 
+## v0.17.0
+- Increased default `monitor_max_run_time_seconds` to 60 seconds.
+- Added `monitor_statement_timeout_seconds` (defaults to 10 seconds) to limit query runtime.
+- Eagerly clear connection pools when a `ActiveRecord::StatementInvalid` or `ActiveRecord::ConnectionTimeoutError`
+  is encountered to attempt to clear bad connections.
+
 ## v0.16.0
 - Add `max_attempts` and `max_run_time` to `Postgres::Vacuum::Jobs::MonitorJob` to avoid backing up the queue. The
   defaults are 1 attempt and 10 seconds, but they can be configured with `monitor_max_attempts` and

README.md

@@ -37,8 +37,10 @@ Postgres::Vacuum::Monitor.configure do |config|
   config.long_running_transaction_threshold_seconds = 10 * 60
   # Optionally change `max_attempts` of the monitor job (default 1)
   config.monitor_max_attempts = 3
-  # Optionally change `max_run_time` of the monitor job (default 10 seconds)
+  # Optionally change `max_run_time` of the monitor job (default 60 seconds)
   config.monitor_max_run_time_seconds = 5
+  # Optionally change the statement timeout of queries (default 10 seconds)
+  config.monitor_statement_timeout_seconds = 5
 end
 ```
 

lib/postgres/vacuum/configuration.rb

@@ -4,19 +4,22 @@ module Postgres
   module Vacuum
     class Configuration
       DEFAULT_LONG_RUNNING_TRANSACTION_THRESHOLD_SECONDS = 5 * 60
-      DEFAULT_MONITOR_MAX_RUN_TIME_SECONDS = 10
+      DEFAULT_MONITOR_MAX_RUN_TIME_SECONDS = 60
       DEFAULT_MONITOR_MAX_ATTEMPTS = 1
+      DEFAULT_MONITOR_STATEMENT_TIMEOUT_SECONDS = 10
 
       attr_accessor :monitor_reporter_class_name,
                     :long_running_transaction_threshold_seconds,
                     :monitor_max_run_time_seconds,
-                    :monitor_max_attempts
+                    :monitor_max_attempts,
+                    :monitor_statement_timeout_seconds
 
       def initialize
         self.monitor_reporter_class_name = nil
         self.long_running_transaction_threshold_seconds = DEFAULT_LONG_RUNNING_TRANSACTION_THRESHOLD_SECONDS
         self.monitor_max_run_time_seconds = DEFAULT_MONITOR_MAX_RUN_TIME_SECONDS
         self.monitor_max_attempts = DEFAULT_MONITOR_MAX_ATTEMPTS
+        self.monitor_statement_timeout_seconds = DEFAULT_MONITOR_STATEMENT_TIMEOUT_SECONDS
       end
     end
   end

lib/postgres/vacuum/jobs/monitor_job.rb

@@ -98,14 +98,42 @@ def with_each_db_name_and_connection
           databases = Set.new
           ActiveRecord::Base.connection_handler.connection_pools.map do |connection_pool|
             db_name = connection_pool.db_config.configuration_hash[:database]
+            next unless databases.add?(db_name)
+
+            # ActiveRecord allocates a connection pool per call to `.establish_connection`
+            # As a result, multiple pools might interact with the same database, so we use
+            # the database name to dedupe.
+            connection_pool.with_connection do |connection|
+              original_timeout = statement_timeout(connection)
+              set_statement_timeout(connection, "#{configured_timeout_seconds}s")
+
+              yield(db_name, connection)
+            ensure
+              set_statement_timeout(connection, original_timeout)
+            end
 
-            # activerecord allocates a connection pool per call to establish_connection
-            # multiple pools might interact with the same database so we use the
-            # database name to dedup
-            connection_pool.with_connection { |conn| yield(db_name, conn) } if databases.add?(db_name)
+          # We want to avoid hanging onto a bad connection that would cause all future
+          # jobs to fail, so we eagerly clear the pool.
+          rescue ActiveRecord::StatementInvalid, ActiveRecord::ConnectionTimeoutError
+            connection_pool.disconnect!
+            raise
           end
         end
 
+        def statement_timeout(connection)
+          result = connection.execute('SHOW statement_timeout').first
+          result['statement_timeout'] if result.present?
+        end
+
+        def set_statement_timeout(connection, timeout)
+          query = ActiveRecord::Base.sanitize_sql(['SET statement_timeout = ?', timeout])
+          connection.execute(query)
+        end
+
+        def configured_timeout_seconds
+          Postgres::Vacuum::Monitor.configuration.monitor_statement_timeout_seconds
+        end
+
         ConfigurationError = Class.new(StandardError)
       end
     end

lib/postgres/vacuum/monitor/version.rb

@@ -3,7 +3,7 @@
 module Postgres
   module Vacuum
     module Monitor
-      VERSION = '0.16.0'
+      VERSION = '0.17.0'
     end
   end
 end

spec/postgres/vacuum/jobs/monitor_job_spec.rb

@@ -168,7 +168,7 @@ def self.report_event(name, attributes = {})
 
         it "reports once for a single database." do
           expect(job.perform).to eq true
-          expect(mock_connection).to have_received(:execute).exactly(5)
+          expect(mock_connection).to have_received(:execute).exactly(8)
         end
 
         context "to different databases" do
@@ -181,10 +181,48 @@ def self.report_event(name, attributes = {})
 
           it "reports twice for two databases" do
             expect(job.perform).to eq true
-            expect(mock_connection).to have_received(:execute).exactly(10)
+            expect(mock_connection).to have_received(:execute).exactly(16)
           end
         end
       end
+
+      describe "statement timeouts" do
+        it "sets and restored statement timeout" do
+          original_timeout = [{ 'statement_timeout' => '20m' }]
+          allow(mock_connection).to receive(:execute).with('SHOW statement_timeout').and_return(original_timeout)
+
+          job.perform
+
+          expect(mock_connection).to have_received(:execute).with("SET statement_timeout = '10s'").ordered
+          expect(mock_connection).to have_received(:execute).with("SET statement_timeout = '20m'").ordered
+        end
+      end
+
+      describe "connection resetting" do
+        it "resets the pool on statement error" do
+          pool = ActiveRecord::Base.connection_handler.connection_pools[0]
+          allow(pool).to receive(:disconnect!)
+
+          error = ActiveRecord::QueryCanceled.new
+          query = Postgres::Vacuum::Monitor::Query.connection_idle_time
+          allow(mock_connection).to receive(:execute).with(query).and_raise(error)
+
+          expect { job.perform }.to raise_error(error)
+          expect(pool).to have_received(:disconnect!)
+        end
+
+        it "resets the pool on connection timeout error" do
+          pool = ActiveRecord::Base.connection_handler.connection_pools[0]
+          allow(pool).to receive(:disconnect!)
+
+          error = ActiveRecord::ConnectionTimeoutError.new
+          query = Postgres::Vacuum::Monitor::Query.connection_idle_time
+          allow(mock_connection).to receive(:execute).with(query).and_raise(error)
+
+          expect { job.perform }.to raise_error(error)
+          expect(pool).to have_received(:disconnect!)
+        end
+      end
     end
   end
 
@@ -216,7 +254,7 @@ def self.report_event(name, attributes = {})
   describe "#max_run_time" do
     context "with default configuration" do
       it "times out after 10 seconds" do
-        expect(job.max_run_time).to eq(10.seconds)
+        expect(job.max_run_time).to eq(60.seconds)
       end
     end