#434

backend/app/services/counts_export_jobs.rb

@@ -4,7 +4,46 @@ require 'fileutils'
 # A durable, bounded export spool. Web and supervised workers share the same POSIX volume.
 class CountsExportJobs
   class Cancelled < StandardError; end
+  class WorkerUnavailable < StandardError; end
   ACTIVE = ['queued', 'running', 'cancelling'].freeze
+  HEARTBEAT_SECONDS = 15
+  QUEUE_SECONDS = CountsExportDownload::MAX_EXPORT_SECONDS + 30
+
+  def self.worker_id
+    @worker_id ||= ".worker-#{ Process.pid }-#{ SecureRandom.hex(8) }"
+  end
+
+  def self.pulse
+    FileUtils.mkdir_p(directory, mode: 0o700)
+    write({ 'id' => worker_id, 'heartbeat_at' => Time.current.to_f })
+  end
+
+  def self.worker_available?
+    Dir.glob(directory.join('.worker-*.json')).any? do |path|
+      JSON.parse(File.read(path)).fetch('heartbeat_at') > Time.current.to_f - HEARTBEAT_SECONDS
+    rescue Errno::ENOENT
+      false
+    end
+  end
+
+  def self.expire_stalled(state)
+    now = Time.current.to_f
+    queued = state['status'] == 'queued'
+    stale = if queued
+              now - state.fetch('created_at', now) > QUEUE_SECONDS ||
+                  (!worker_available? && now - state.fetch('created_at', now) > HEARTBEAT_SECONDS)
+            else
+              ['running', 'cancelling'].include?(state['status']) &&
+                  now - state.fetch('heartbeat_at', now) > HEARTBEAT_SECONDS
+            end
+    return unless stale
+
+    state.merge!('status' => 'failed', 'error' => queued ? 'worker_unavailable' : 'worker_stalled',
+        'updated_at' => now)
+    remove_download(state)
+    state.delete('token')
+    write(state)
+  end
 
   def self.directory = CountsExportDownload.directory.join('jobs')
   def self.verifier = Rails.application.message_verifier('counts_export_job')
@@ -91,10 +130,14 @@ class CountsExportJobs
       payload[time_key] = Time.current.utc.iso8601(6)
     end
     exporter = CountsExporter.new(payload, language_code: language)
+    authority = HistoryAccess.authority
     synchronise do
       clean
+      raise WorkerUnavailable unless worker_available?
       active = Dir.glob(directory.join('*.json')).count do |path|
-        ACTIVE.include?(JSON.parse(File.read(path))['status'])
+        existing = JSON.parse(File.read(path))
+        expire_stalled(existing)
+        ACTIVE.include?(existing['status'])
       end
       raise CountsExportDownload::Busy if active >= CountsExportDownload::MAX_CONCURRENT_EXPORTS
 
@@ -102,8 +145,10 @@ class CountsExportJobs
       state = { 'id' => id, 'owner' => owner, 'payload' => payload,
           'language' => language, 'resource' => exporter.resource,
           'format' => exporter.format, 'status' => 'queued',
+          'phase' => 'queued', 'created_at' => Time.current.to_f,
+          'updated_at' => Time.current.to_f, 'heartbeat_at' => Time.current.to_f,
           'processed' => 0, 'total' => nil,
-          **HistoryAccess.authority,
+          **authority,
           'expires_at' => (Time.current + CountsExportDownload::EXPIRY).to_f }
       write(state)
       verifier.generate({ 'id' => id, 'owner' => owner },
@@ -122,12 +167,15 @@ class CountsExportJobs
   def self.status(token, owner, cancel: false)
     id = authorised_id(token, owner)
     return unless id
+    snapshot = read(id)
+    return unless snapshot && snapshot['owner'] == owner
+    return unless HistoryAccess.authorised_artifact?(snapshot)
 
     synchronise do
       state = read(id)
       return unless state && state['owner'] == owner && state['expires_at'] > Time.current.to_f
 
-      return unless HistoryAccess.authorised_artifact?(state)
+      expire_stalled(state)
 
       if cancel && ['queued', 'running', 'ready'].include?(state['status'])
         state['status'] = state['status'] == 'running' ? 'cancelling' : 'cancelled'
@@ -135,24 +183,28 @@ class CountsExportJobs
         state.delete('token')
         write(state)
       end
-      state.slice('status', 'processed', 'total', 'format', 'token', 'error')
+      state.slice('status', 'phase', 'processed', 'total', 'format', 'token', 'error',
+          'updated_at', 'heartbeat_at')
     end
   end
 
-  def self.checkpoint(id, processed = nil, total = nil)
+  def self.checkpoint(id, processed = nil, total = nil, phase: nil)
     synchronise do
       state = read(id)
       unless state && state['status'] == 'running' && state['expires_at'] > Time.current.to_f
         raise Cancelled
       end
-      if processed
-        state.merge!('processed' => processed, 'total' => total)
+      if processed || phase
+        state.merge!('processed' => processed, 'total' => total) unless processed.nil?
+        state['phase'] = phase if phase
+        state['updated_at'] = Time.current.to_f
         write(state)
       end
     end
   end
 
   def self.work_one
+    pulse
     FileUtils.mkdir_p(directory, mode: 0o700)
     Dir.glob(directory.join('*.json')).sort.each do |path|
       id = File.basename(path, '.json')
@@ -171,12 +223,15 @@ class CountsExportJobs
             next
           end
           candidate['status'] = 'running'
+          candidate['phase'] = 'aggregating'
+          candidate['heartbeat_at'] = Time.current.to_f
+          candidate['updated_at'] = Time.current.to_f
           write(candidate)
           candidate
         end
         next unless state
 
-        perform(state)
+        supervise(state)
         return true
       end
     end
@@ -184,6 +239,80 @@ class CountsExportJobs
     false
   end
 
+  def self.supervise(state)
+    id = state.fetch('id')
+    ApplicationRecord.connection_handler.clear_all_connections!
+    pid = Process.fork do
+      Signal.trap('TERM') { raise Cancelled }
+      Signal.trap('INT') { raise Cancelled }
+      begin
+        perform(state)
+      ensure
+        ApplicationRecord.connection_handler.clear_all_connections!
+        Process.exit!(0)
+      end
+    end
+    started = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+    termination = nil
+    reaped = false
+    begin
+      loop do
+        if Process.waitpid(pid, Process::WNOHANG)
+          reaped = true
+          break
+        end
+        pulse
+        now = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+        terminal = synchronise do
+          current = read(id)
+          if current && current['status'] == 'running'
+            if now - started > CountsExportDownload::MAX_EXPORT_SECONDS
+              current.merge!('status' => 'failed', 'error' => 'generation_timeout')
+            end
+            current['heartbeat_at'] = Time.current.to_f
+            write(current)
+          end
+          !current || !['running', 'ready'].include?(current['status'])
+        end
+        begin
+          if terminal && termination.nil?
+            Process.kill('TERM', pid)
+            termination = now
+          elsif termination && now - termination > 5
+            Process.kill('KILL', pid)
+          end
+        rescue Errno::ESRCH
+          # The child may exit between the nonblocking wait and signal delivery.
+          nil
+        end
+        sleep(1)
+      end
+    ensure
+      unless reaped
+        begin
+          Process.kill('KILL', pid)
+        rescue Errno::ESRCH
+          nil
+        end
+        Process.waitpid(pid)
+      end
+      synchronise do
+        current = read(id)
+        if current && current['status'] != 'ready'
+          if ['running', 'cancelling'].include?(current['status'])
+            current['status'] = current['status'] == 'cancelling' ? 'cancelled' : 'failed'
+            current['error'] = 'interrupted' if current['status'] == 'failed'
+          end
+          remove_download(current)
+          remove_partial(id)
+          current.delete('token')
+          current['updated_at'] = Time.current.to_f
+          write(current)
+        end
+      end
+    end
+  end
+
   def self.perform(state)
     actor = state['actor_id'] && User.find_by(id: state['actor_id'])
     actor = nil unless state['admin'] && actor&.admin? && !actor.banned?
@@ -221,18 +350,21 @@ class CountsExportJobs
           last = now
         end
       end
+      phase = ->(value) { checkpoint(id, phase: value) }
       token = CountsExportDownload.prepare(exporter, state.fetch('owner'),
-          progress:, check:, identifier: id)
+          progress:, phase:, check:, identifier: id)
+      file = CountsExportDownload.find(token, state.fetch('owner'))
+      raise Cancelled unless file
+
       synchronise do
         current = read(id)
         if (!current || current['status'] != 'running' ||
             current['expires_at'] <= Time.current.to_f)
-          file = CountsExportDownload.find(token, state.fetch('owner'))
-          File.unlink(file.fetch('path')) if file
+          File.unlink(file.fetch('path')) if File.exist?(file.fetch('path'))
           raise Cancelled
         end
-        file = CountsExportDownload.find(token, state.fetch('owner'))
         current.merge!('status' => 'ready', 'token' => token,
+            'updated_at' => Time.current.to_f,
             'download_file' => File.basename(file.fetch('path')))
         current['total'] = current['processed']
         write(current)
@@ -242,8 +374,11 @@ class CountsExportJobs
       synchronise do
         current = read(id)
         if current
-          current['status'] = error.is_a?(Cancelled) ? 'cancelled' : 'failed'
-          current['error'] = error.class.name
+          unless current['status'] == 'failed'
+            current['status'] = error.is_a?(Cancelled) ? 'cancelled' : 'failed'
+            current['error'] = error.class.name
+          end
+          current['updated_at'] = Time.current.to_f
           write(current)
         end
       end