109 lines
2.6 KiB
Ruby
109 lines
2.6 KiB
Ruby
# typed: strict
|
|
module Tasks
|
|
class InterruptableTask
|
|
extend T::Sig
|
|
extend T::Helpers
|
|
abstract!
|
|
PB_FORMAT = T.let("%B %c/%C (%r/sec) %J%% %a %E", String)
|
|
|
|
sig { params(log_sink: T.any(IO, StringIO)).void }
|
|
def initialize(log_sink: $stderr)
|
|
@interrupt_monitor =
|
|
T.let(
|
|
Tasks::UserInterruptMonitor.new(log_sink:),
|
|
Tasks::UserInterruptMonitor,
|
|
)
|
|
@log_sink = log_sink
|
|
log("Press Ctrl+C to cleanly interrupt and save progress")
|
|
end
|
|
|
|
sig { abstract.returns(String) }
|
|
def progress_key
|
|
end
|
|
|
|
sig { params(start_at: T.nilable(String)).returns(T.nilable(String)) }
|
|
def get_progress(start_at)
|
|
value =
|
|
if start_at == "last"
|
|
value = GlobalState.get(progress_key)
|
|
else
|
|
start_at
|
|
end
|
|
|
|
if value
|
|
log(
|
|
"[progress key: #{progress_key.bold}] [resuming from: #{value.to_s.bold}]",
|
|
)
|
|
else
|
|
log("[progress key: #{progress_key.bold}] [no saved progress]")
|
|
end
|
|
|
|
value
|
|
end
|
|
|
|
sig { params(value: String).void }
|
|
def save_progress(value)
|
|
GlobalState.set(progress_key, value)
|
|
end
|
|
|
|
sig { params(message: String).void }
|
|
def log(message)
|
|
@log_sink.puts(message)
|
|
end
|
|
|
|
sig { returns(T::Boolean) }
|
|
def interrupted?
|
|
@interrupt_monitor.interrupted?
|
|
end
|
|
|
|
sig { params(total: T.nilable(Integer)).returns(ProgressBar::Base) }
|
|
def create_progress_bar(total)
|
|
ProgressBar.create(
|
|
total: total,
|
|
progress_mark: " ",
|
|
remainder_mark: " ",
|
|
format: PB_FORMAT,
|
|
)
|
|
end
|
|
|
|
sig { void }
|
|
def run
|
|
begin
|
|
run_impl
|
|
ensure
|
|
end_profiling!
|
|
end
|
|
end
|
|
|
|
sig { abstract.void }
|
|
def run_impl
|
|
end
|
|
|
|
sig { void }
|
|
def start_profiling!
|
|
return unless ENV["PROFILE"]
|
|
@log_sink.puts "starting profiling"
|
|
RubyProf.start
|
|
end
|
|
|
|
sig { void }
|
|
def end_profiling!
|
|
return unless ENV["PROFILE"]
|
|
return unless RubyProf.running?
|
|
base = "profiler/#{progress_key}"
|
|
FileUtils.mkdir_p(base) unless File.exist?(base)
|
|
result = RubyProf.stop
|
|
File.open("#{base}/profile.txt", "w") do |f|
|
|
RubyProf::GraphPrinter.new(result).print(f, { min_percent: 1 })
|
|
end
|
|
File.open("#{base}/profile.html", "w") do |f|
|
|
RubyProf::CallStackPrinter.new(result).print(f, { min_percent: 1 })
|
|
end
|
|
File.open("#{base}/profile.rubyprof", "w") do |f|
|
|
RubyProf::SpeedscopePrinter.new(result).print(f, { min_percent: 1 })
|
|
end
|
|
@log_sink.puts "profiling results saved to #{base}"
|
|
end
|
|
end
|
|
end
|