391 行
13 KiB
Ruby
391 行
13 KiB
Ruby
require 'json'
|
|
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')
|
|
|
|
def self.synchronise
|
|
FileUtils.mkdir_p(directory, mode: 0o700)
|
|
File.open(directory.join('.registry.lock'), File::RDWR | File::CREAT, 0o600) do |lock|
|
|
lock.flock(File::LOCK_EX)
|
|
yield
|
|
end
|
|
end
|
|
|
|
def self.write(state)
|
|
path = directory.join("#{ state.fetch('id') }.json")
|
|
temporary = "#{ path }.#{ SecureRandom.hex(8) }.tmp"
|
|
File.open(temporary, 'w', 0o600) do |file|
|
|
file.write(JSON.generate(state))
|
|
file.flush
|
|
file.fsync
|
|
end
|
|
File.rename(temporary, path)
|
|
ensure
|
|
File.unlink(temporary) if temporary && File.exist?(temporary)
|
|
end
|
|
|
|
def self.read(id)
|
|
JSON.parse(File.read(directory.join("#{ id }.json")))
|
|
rescue Errno::ENOENT
|
|
nil
|
|
end
|
|
|
|
def self.remove_download(state)
|
|
name = state['download_file']
|
|
return unless name&.match?(/\Acounts-[a-zA-Z0-9_-]+\.export\z/)
|
|
|
|
File.unlink(CountsExportDownload.directory.join(name))
|
|
rescue Errno::ENOENT
|
|
nil
|
|
end
|
|
|
|
def self.remove_partial(id)
|
|
Dir.glob(CountsExportDownload.directory.join("counts-#{ id }-*")).each do |path|
|
|
File.unlink(path)
|
|
rescue Errno::ENOENT
|
|
nil
|
|
end
|
|
end
|
|
|
|
def self.clean
|
|
Dir.glob(directory.join('*.tmp')).each do |path|
|
|
File.unlink(path) if File.mtime(path) < Time.current - CountsExportDownload::EXPIRY
|
|
end
|
|
Dir.glob(directory.join('*.run')).each do |path|
|
|
id = File.basename(path, '.run')
|
|
next if directory.join("#{ id }.json").exist?
|
|
|
|
File.open(path, File::RDWR) do |lock|
|
|
File.unlink(path) if lock.flock(File::LOCK_EX | File::LOCK_NB)
|
|
end
|
|
end
|
|
Dir.glob(directory.join('*.json')).each do |path|
|
|
state = JSON.parse(File.read(path))
|
|
next if state.fetch('expires_at') > Time.current.to_f
|
|
|
|
run = directory.join("#{ state.fetch('id') }.run")
|
|
File.open(run, File::RDWR | File::CREAT, 0o600) do |lock|
|
|
unless lock.flock(File::LOCK_EX | File::LOCK_NB)
|
|
state['status'] = 'cancelling'
|
|
write(state)
|
|
next
|
|
end
|
|
remove_download(state)
|
|
remove_partial(state.fetch('id'))
|
|
File.unlink(path)
|
|
File.unlink(run)
|
|
end
|
|
end
|
|
end
|
|
|
|
def self.start(payload, owner, language)
|
|
payload = payload.dup
|
|
time_key = payload.key?('definition') ? 'execution_at' : 'at'
|
|
if payload[time_key].nil? || payload[time_key] == ''
|
|
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|
|
|
existing = JSON.parse(File.read(path))
|
|
expire_stalled(existing)
|
|
ACTIVE.include?(existing['status'])
|
|
end
|
|
raise CountsExportDownload::Busy if active >= CountsExportDownload::MAX_CONCURRENT_EXPORTS
|
|
|
|
id = SecureRandom.hex(24)
|
|
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,
|
|
**authority,
|
|
'expires_at' => (Time.current + CountsExportDownload::EXPIRY).to_f }
|
|
write(state)
|
|
verifier.generate({ 'id' => id, 'owner' => owner },
|
|
expires_in: CountsExportDownload::EXPIRY, purpose: 'counts_export_job')
|
|
end
|
|
end
|
|
|
|
def self.authorised_id(token, owner)
|
|
value = verifier.verified(token, purpose: 'counts_export_job')
|
|
return unless value.is_a?(Hash) && value['owner'] == owner
|
|
return unless value['id']&.match?(/\A[a-f0-9]{48}\z/)
|
|
|
|
value['id']
|
|
end
|
|
|
|
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
|
|
|
|
expire_stalled(state)
|
|
|
|
if cancel && ['queued', 'running', 'ready'].include?(state['status'])
|
|
state['status'] = state['status'] == 'running' ? 'cancelling' : 'cancelled'
|
|
remove_download(state)
|
|
state.delete('token')
|
|
write(state)
|
|
end
|
|
state.slice('status', 'phase', 'processed', 'total', 'format', 'token', 'error',
|
|
'updated_at', 'heartbeat_at')
|
|
end
|
|
end
|
|
|
|
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 || 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')
|
|
File.open(directory.join("#{ id }.run"), File::RDWR | File::CREAT, 0o600) do |lock|
|
|
next unless lock.flock(File::LOCK_EX | File::LOCK_NB)
|
|
|
|
state = synchronise do
|
|
candidate = read(id)
|
|
next unless candidate && ACTIVE.include?(candidate['status'])
|
|
|
|
if candidate['status'] != 'queued'
|
|
candidate['status'] = candidate['status'] == 'cancelling' ? 'cancelled' : 'failed'
|
|
remove_partial(id)
|
|
candidate['error'] = 'interrupted'
|
|
write(candidate)
|
|
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
|
|
|
|
supervise(state)
|
|
return true
|
|
end
|
|
end
|
|
synchronise { clean }
|
|
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?
|
|
HistoryAccess.set(user: actor) do
|
|
unless HistoryAccess.authorised_artifact?(state)
|
|
synchronise do
|
|
current = read(state.fetch('id'))
|
|
current['status'] = 'failed'
|
|
write(current)
|
|
end
|
|
return
|
|
end
|
|
perform_authorised(state)
|
|
end
|
|
end
|
|
|
|
def self.perform_authorised(state)
|
|
id = state.fetch('id')
|
|
token = nil
|
|
begin
|
|
checkpoint(id)
|
|
exporter = CountsExporter.new(state.fetch('payload'), language_code: state.fetch('language'))
|
|
last = Process.clock_gettime(Process::CLOCK_MONOTONIC)
|
|
progress = lambda do |processed, total|
|
|
now = Process.clock_gettime(Process::CLOCK_MONOTONIC)
|
|
if processed.zero? || now - last >= 0.25 || processed == total
|
|
checkpoint(id, processed, total)
|
|
last = now
|
|
end
|
|
end
|
|
check = lambda do
|
|
now = Process.clock_gettime(Process::CLOCK_MONOTONIC)
|
|
if now - last >= 0.25
|
|
checkpoint(id)
|
|
last = now
|
|
end
|
|
end
|
|
phase = ->(value) { checkpoint(id, phase: value) }
|
|
token = CountsExportDownload.prepare(exporter, state.fetch('owner'),
|
|
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.unlink(file.fetch('path')) if File.exist?(file.fetch('path'))
|
|
raise Cancelled
|
|
end
|
|
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)
|
|
end
|
|
rescue StandardError => error
|
|
remove_partial(id)
|
|
synchronise do
|
|
current = read(id)
|
|
if current
|
|
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
|
|
unless error.is_a?(Cancelled)
|
|
Rails.logger.error("counts_export_job #{ id }: #{ error.class }")
|
|
end
|
|
end
|
|
end
|
|
end
|