Add pg_notify support to reduce polling latency

README.md

@@ -112,6 +112,12 @@ To determine what happens when a job is queued, you can set Que's mode. There ar
 
 **If you're using ActiveRecord to dump your database's schema, [set your schema_format to :sql](http://guides.rubyonrails.org/migrations.html#types-of-schema-dumps) so that Que's table structure is managed correctly.** (You can use schema_format as :ruby if you want but keep in mind this is highly advised against, as some parts of Que will not work.)
 
+### Reducing polling latency with pg_notify
+
+By default, idle workers only notice new jobs on their next poll (`--wake-interval`/`wake_interval`, 5 seconds by default). Running `Que.migrate!` up to the latest schema version adds a trigger that calls `pg_notify` whenever a job becomes available, and passing `--listen` to `bin/que` (or `listen: true` to `Que::WorkerGroup.start`) makes idle workers wake up on that notification instead of waiting out the full interval. Polling itself is unaffected and still runs as normal - notifications are purely a way to shorten the wait, so a missed one just means the job is picked up on the next poll as before.
+
+This holds one dedicated connection open for the lifetime of each worker group to listen on, in addition to the connections workers use to run queries. If you're using a pooled adapter (ActiveRecord, Sequel, Pond), make sure your pool has capacity for one extra connection per worker group when enabling `--listen`.
+
 
 ## Related Projects
 

bin/que

@@ -51,6 +51,10 @@ OptionParser.new do |opts|
     options.run_at_cursor = true
   end
 
+  opts.on("--listen", "Listen for pg_notify() wakeups so new jobs are picked up faster than wake-interval (default: false)") do
+    options.listen = true
+  end
+
   opts.on("-c", "--worker-count [COUNT]", Integer, "Start this many threaded workers") do |worker_count|
     options.worker_count = worker_count
   end
@@ -108,6 +112,7 @@ run_at_cursor    = options.run_at_cursor || false
 worker_count     = options.worker_count  || 1
 timeout          = options.timeout
 secondary_queues = options.secondary_queues || []
+listen           = options.listen || false
 
 Que.logger ||= Logger.new($stdout)
 
@@ -134,6 +139,7 @@ worker_group = Que::WorkerGroup.start(
   lock_budget: options.lock_budget,
   secondary_queues: secondary_queues,
   run_at_cursor: run_at_cursor,
+  listen: listen,
 )
 
 if options.metrics_port

lib/que.rb

@@ -6,6 +6,7 @@
 require_relative "que/job"
 require_relative "que/job_timeout_error"
 require_relative "que/leaky_bucket"
+require_relative "que/listener"
 require_relative "que/locker"
 require_relative "que/migrations"
 require_relative "que/sql"

lib/que/listener.rb

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+module Que
+  # Listens for pg_notify() messages announcing that a job has become
+  # available on one of the given queues, and uses them to wake up idle
+  # workers early instead of leaving them to wait out their next poll.
+  #
+  # This is a supplement to polling, not a replacement for it - workers still
+  # poll on every cycle regardless of whether they were woken by a
+  # notification or by their poll interval elapsing. That means a missed or
+  # delayed notification (the listener hasn't (re)connected yet, a
+  # notification arrives while we're reconnecting, etc) only costs the
+  # latency of the next poll rather than losing the job.
+  #
+  # Requires a dedicated connection for its entire lifetime (to hold the
+  # LISTEN), which for pooled adapters (ActiveRecord/Sequel/Pond) means the
+  # pool needs capacity for one extra connection per Listener on top of
+  # however many workers are checking connections out to run queries.
+  class Listener
+    RECONNECT_WAIT = 1.0 # seconds
+    POLL_TIMEOUT = 1.0 # seconds; bounds how quickly #stop is noticed
+
+    def self.channel_for_queue(queue)
+      "que_job_available_#{queue}"
+    end
+
+    def initialize(queues:, wakeup:)
+      @queues = queues
+      @wakeup = wakeup
+      @stop = false
+    end
+
+    def start
+      @thread = Thread.new { run }
+      @thread.name = "que-listener"
+      self
+    end
+
+    def stop
+      @stop = true
+      @thread&.join
+    end
+
+    private
+
+    def run
+      until @stop
+        begin
+          listen_and_wait
+        rescue PG::Error, Adapters::UnavailableConnection => e
+          Que.logger&.warn(
+            event: "que.listener.connection_error",
+            msg: "Listener lost its connection, reconnecting",
+            error: e.inspect,
+          )
+          sleep(RECONNECT_WAIT) unless @stop
+        end
+      end
+    end
+
+    def listen_and_wait
+      Que.adapter.checkout do |conn|
+        listen(conn)
+
+        begin
+          wait_for_notifications(conn) until @stop
+        ensure
+          unlisten(conn)
+        end
+      end
+    end
+
+    def listen(conn)
+      @queues.each do |queue|
+        conn.exec(%(LISTEN #{conn.quote_ident(self.class.channel_for_queue(queue))}))
+      end
+    end
+
+    def unlisten(conn)
+      conn.exec("UNLISTEN *")
+      nil while conn.notifies
+    end
+
+    def wait_for_notifications(conn)
+      conn.wait_for_notify(POLL_TIMEOUT) { wake! }
+    end
+
+    def wake!
+      @wakeup.fetch(:mutex).synchronize { @wakeup.fetch(:condvar).broadcast }
+    end
+  end
+end

lib/que/migrations.rb

@@ -5,7 +5,7 @@ module Migrations
     # In order to ship a schema change, add the relevant up and down sql files
     # to the migrations directory, and bump the version both here and in the
     # add_que generator template.
-    CURRENT_VERSION = 6
+    CURRENT_VERSION = 7
 
     class << self
       def migrate!(options = { version: CURRENT_VERSION })

lib/que/migrations/7/down.sql

@@ -0,0 +1,6 @@
+BEGIN;
+
+DROP TRIGGER que_job_notify ON que_jobs;
+DROP FUNCTION que_job_notify();
+
+COMMIT;

lib/que/migrations/7/up.sql

@@ -0,0 +1,44 @@
+BEGIN;
+
+CREATE FUNCTION que_job_notify() RETURNS trigger AS $$
+  DECLARE
+    payload json;
+  BEGIN
+    -- Don't bother notifying about jobs that aren't ready to run yet - a
+    -- worker polling for them later will pick them up regardless.
+    IF NEW.run_at IS NOT NULL AND NEW.run_at > now() THEN
+      RETURN null;
+    END IF;
+
+    -- There's a size limit to what can be broadcast via LISTEN/NOTIFY, so
+    -- rather than throw errors when someone enqueues a big job, just
+    -- broadcast the most pertinent information, and let the worker query for
+    -- the record after it's taken the lock, exactly as it does when polling.
+    SELECT row_to_json(t)
+    INTO payload
+    FROM (
+      SELECT
+        NEW.queue    AS queue,
+        NEW.priority AS priority,
+        NEW.job_id   AS job_id,
+        to_char(NEW.run_at AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.US"Z"') AS run_at
+    ) t;
+
+    -- One channel per queue, so listeners only subscribe to the queues they
+    -- actually work.
+    PERFORM pg_notify('que_job_available_' || NEW.queue, payload::text);
+
+    RETURN null;
+  END
+$$
+LANGUAGE plpgsql;
+
+-- Allow bulk-enqueue code paths to skip this trigger for performance by
+-- running `SET LOCAL que.skip_notify = 'true'` in their transaction.
+CREATE TRIGGER que_job_notify
+  AFTER INSERT ON que_jobs
+  FOR EACH ROW
+  WHEN (NOT coalesce(current_setting('que.skip_notify', true), '') = 'true')
+  EXECUTE PROCEDURE que_job_notify();
+
+COMMIT;

lib/que/worker.rb

@@ -122,10 +122,12 @@ def initialize(
       lock_window: nil,
       lock_budget: nil,
       secondary_queues: [],
-      run_at_cursor: false
+      run_at_cursor: false,
+      wakeup: nil
     )
       @queue = queue
       @wake_interval = wake_interval
+      @wakeup = wakeup
       @tracer = LongRunningMetricTracer.new(self)
       @stop = false # instruct worker to stop
       @stopped = false # mark worker as having stopped
@@ -157,7 +159,7 @@ def work_loop
           when :job_not_found
             Que.logger&.debug(event: "que.job_not_found", wake_interval: @wake_interval)
             @tracer.trace(SleepingSecondsTotal, queue: @queue, primary_queue: @queue) do
-              sleep(@wake_interval)
+              wait_for_wakeup_or_timeout
             end
           when :job_worked
             nil # immediately find a new job to work
@@ -290,6 +292,17 @@ def collect_metrics
 
     private
 
+    # Wait up to wake_interval seconds before polling again. If we've been given a
+    # wakeup mutex/condvar (i.e. pg_notify listening is enabled), wake up as soon as
+    # a job becomes available rather than always waiting out the full interval - the
+    # bounded wait still falls back to plain polling if no notification arrives.
+    def wait_for_wakeup_or_timeout
+      return sleep(@wake_interval) unless @wakeup
+
+      mutex = @wakeup.fetch(:mutex)
+      mutex.synchronize { @wakeup.fetch(:condvar).wait(mutex, @wake_interval) }
+    end
+
     # Set the error and retry with back-off
     def handle_job_failure(error, job)
       count = job[:error_count].to_i + 1

lib/que/worker_group.rb

@@ -17,16 +17,20 @@ class WorkerGroup
       ),
     ].freeze
 
-    def self.start(count, **kwargs)
+    def self.start(count, listen: false, **kwargs)
       Que.logger.info(
         msg: "Starting workers",
         event: "que.worker.start",
         worker_count: count,
       )
 
-      workers = Array.new(count) { Worker.new(**kwargs) }
+      wakeup = { mutex: Mutex.new, condvar: ConditionVariable.new } if listen
+
+      workers = Array.new(count) { Worker.new(wakeup: wakeup, **kwargs) }
       worker_threads = workers.map { |worker| Thread.new { worker.work_loop } }
 
+      listener = start_listener(kwargs, wakeup) if listen
+
       # We want to ensure that our control flow for Que threads and our metric endpoints
       # are prioritised above CPU intensive work of jobs. If people use the metrics
       # endpoint as a liveness check then its essential we don't timeout when we do busy
@@ -40,12 +44,22 @@ def self.start(count, **kwargs)
         thread.name = "worker-thread-#{index}"
       end
 
-      new(workers, worker_threads)
+      new(workers, worker_threads, listener, wakeup)
+    end
+
+    def self.start_listener(kwargs, wakeup)
+      queue = kwargs[:queue] || Worker::DEFAULT_QUEUE
+      queues = Array(queue).concat(Array(kwargs[:secondary_queues])).uniq
+
+      Listener.new(queues: queues, wakeup: wakeup).start
     end
+    private_class_method :start_listener
 
-    def initialize(workers, worker_threads)
+    def initialize(workers, worker_threads, listener = nil, wakeup = nil)
       @workers = workers
       @worker_threads = worker_threads
+      @listener = listener
+      @wakeup = wakeup
     end
 
     attr_reader :workers
@@ -69,6 +83,10 @@ def stop(timeout = DEFAULT_STOP_TIMEOUT)
 
       @workers.each(&:stop!)
 
+      # If workers are parked waiting on a notification (rather than plain polling),
+      # wake them immediately so they notice @stop without waiting out wake_interval.
+      wake_all_workers!
+
       # Asynchronously stop all workers, sending a join method call with a timeout. This
       # is done asynchronously to ensure we start the timeout at the same time for each
       # worker, instead of applying them cumulatively.
@@ -102,7 +120,17 @@ def stop(timeout = DEFAULT_STOP_TIMEOUT)
       # exception against the job it was working.
       @worker_threads.each(&:join)
 
+      @listener&.stop
+
       Que.logger.info(msg: "All workers have finished", event: "que.worker.finish")
     end
+
+    private
+
+    def wake_all_workers!
+      return unless @wakeup
+
+      @wakeup.fetch(:mutex).synchronize { @wakeup.fetch(:condvar).broadcast }
+    end
   end
 end

spec/integration/integration_spec.rb

@@ -103,6 +103,24 @@
     end
   end
 
+  context "with pg_notify listening enabled" do
+    it "works a job enqueued while idle well before the wake interval elapses" do
+      start = nil
+
+      with_workers(1, wake_interval: 20, listen: true) do
+        sleep 0.2 # give the listener time to issue LISTEN before we enqueue
+
+        start = Time.now
+        FakeJob.enqueue(1)
+        wait_for_jobs_to_be_worked(timeout: 2)
+      end
+
+      expect(Time.now - start).to be < 2
+      expect(QueJob.count).to eq(0)
+      expect(FakeJob.log).to eq([1])
+    end
+  end
+
   context "with jobs that exceed stop timeout" do
     it "raises Que::JobTimeoutError" do
       SleepJob.enqueue(5) # sleep 5s