task_progress 0.0.1

data/lib/task_progress/rails_defaults.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+# rbs_inline: enabled
+
+class TaskProgress
+  # What TaskProgress.broadcast and the `task_progress_tree` helper use when Rails is there and
+  # nothing else is given: the record's dom_id, Turbo::StreamsChannel, Rails.cache, and tokens
+  # signed with a key the app's secret_key_base derives. Without Rails, nothing - and the caller
+  # says what to broadcast to. The record is streamed to as `[record, :task_progress]`
+  module RailsDefaults # :nodoc:
+    class << self
+      # The id the tree's element has, for a record or a string
+      #: (untyped) -> String
+      def id(target)
+        return target.to_s if target.is_a?(String) || target.is_a?(Symbol) || !rails?
+        ActionView::RecordIdentifier.dom_id(target, :task_progress) # steep:ignore
+      end
+
+      # The store key the latest frame is kept under
+      #: (untyped) -> String
+      def key(target) = "task_progress/#{id(target)}"
+
+      # The defaults for what is not given, or nothing outside Rails
+      #: (untyped) -> Hash[Symbol, untyped]
+      def for(target)
+        return {} unless rails?
+        defaults = {store: ::Rails.cache, tokens:, cancel_url: cancellations_path} # steep:ignore
+        return defaults unless defined?(::Turbo::StreamsChannel)
+        defaults.merge(to: ->(html) { ::Turbo::StreamsChannel.broadcast_stream_to(target, :task_progress, content: html) }) # steep:ignore
+      end
+
+      # Tokens for the cancel buttons, signed with a key derived from the app's secret
+      #: () -> Token
+      def tokens
+        Token.new(::Rails.application.key_generator.generate_key("task_progress cancellations")) # steep:ignore
+      end
+
+      # Where the cancel buttons post, or nil when the app has turned that off
+      #: () -> String?
+      def cancellations_path = (::Rails.application.config.task_progress.cancellations_path if rails?) # steep:ignore
+
+      # Whether this runs in a Rails app
+      #: () -> bool
+      def rails? = !!(defined?(::Rails.application) && ::Rails.application) # steep:ignore
+    end
+  end
+end

data/lib/task_progress/railtie.rb

@@ -0,0 +1,91 @@
+# frozen_string_literal: true
+
+class TaskProgress
+  # Loaded only inside a Rails app: trees as job arguments (see JobSerializer), the
+  # `task_progress_tree` helper in every view, the browser client at `/task_progress/js` (and on the
+  # asset path, and pinned in the import map when importmap-rails is in use), a route for the
+  # cancel buttons - at `/task_progress/cancellations`, unless `config.task_progress.cancellations_path`
+  # says elsewhere, or is nil for none - and `rails task_progress:vendor`
+  class Railtie < ::Rails::Railtie
+    config.task_progress = ActiveSupport::OrderedOptions.new
+    config.task_progress.cancellations_path = "/task_progress/cancellations"
+    # The browser client, as it is - an asset pipeline would fingerprint the file you import, and
+    # leave the relative imports inside it pointing at names it does not serve
+    config.task_progress.js_path = "/task_progress/js"
+    # Pins `task_progress`, `task_progress/react` and the rest to it, when importmap-rails is in use
+    config.task_progress.importmap = true
+    # Where `rails task_progress:vendor` copies the client to, for a bundler to take it from
+    config.task_progress.vendor_path = "vendor/javascript/task_progress"
+    # Where pages get handed-off trees from, by the links TaskProgress.for_page hands out
+    config.task_progress.trees_path = "/task_progress/trees"
+    # Where handed-off trees are kept - a TaskProgress::RedisStreamStore, say. Unset, see TaskProgress.stream_store
+    config.task_progress.stream_store = nil
+
+    initializer "task_progress.config" do |app|
+      TaskProgress.stream_store = app.config.task_progress.stream_store if app.config.task_progress.stream_store
+      TaskProgress.trees_path = app.config.task_progress.trees_path if app.config.task_progress.trees_path
+    end
+
+    config.after_initialize { TaskProgress.tokens = TaskProgress::RailsDefaults.tokens }
+
+    initializer "task_progress.jobs" do
+      ActiveSupport.on_load(:active_job) { TaskProgress::JobSerializer.install }
+    end
+
+    initializer "task_progress.helper" do
+      ActiveSupport.on_load(:action_view) { include TaskProgress::TurboHelper }
+    end
+
+    initializer "task_progress.assets" do |app|
+      app.config.assets.paths << TaskProgress.js_path if app.config.respond_to?(:assets)
+    end
+
+    # Before importmap-rails draws the app's own pins, so that a pin of the same name there wins
+    initializer "task_progress.importmap", before: "importmap" do |app|
+      options = app.config.task_progress
+      if app.config.respond_to?(:importmap) && options.importmap && options.js_path
+        app.config.importmap.paths << File.expand_path("../../config/importmap.rb", __dir__)
+      end
+    end
+
+    # A vendored copy is not updated with the gem - say so, rather than serve the old one
+    config.after_initialize do |app|
+      vendored = TaskProgress.vendored_js_version(app.root.join(app.config.task_progress.vendor_path).to_s)
+      if vendored && vendored != TaskProgress::VERSION
+        Rails.logger&.warn("task_progress #{TaskProgress::VERSION} is installed, but #{app.config.task_progress.vendor_path} has the client from #{vendored} - run `rails task_progress:vendor`")
+      end
+    end
+
+    rake_tasks do
+      namespace :task_progress do
+        desc "Copy the browser client into config.task_progress.vendor_path, or DIR"
+        task :vendor do
+          dir = ENV["DIR"] || Rails.root.join(Rails.application.config.task_progress.vendor_path).to_s
+          TaskProgress.vendor_js(dir)
+          puts "The task_progress client #{TaskProgress::VERSION} is in #{dir}"
+        end
+      end
+    end
+
+    initializer "task_progress.routes" do |app|
+      path = app.config.task_progress.cancellations_path
+      if path
+        app.routes.append do
+          mount TaskProgress::CancellationsApp.new(store: Rails.cache, tokens: TaskProgress::RailsDefaults.tokens) => path
+        end
+      end
+      trees = app.config.task_progress.trees_path
+      if trees
+        app.routes.append do
+          mount TaskProgress::TreesApp.new(tokens: TaskProgress::RailsDefaults.tokens) => trees
+        end
+      end
+      js = app.config.task_progress.js_path
+      if js
+        app.routes.append do
+          mount Rack::Files.new(TaskProgress.js_path, {"cache-control" => "no-cache"}) => js
+        end
+      end
+    end
+  end
+end

data/lib/task_progress/recording_output.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+# rbs_inline: enabled
+
+class TaskProgress
+  # Keeps every batch it gets, for a test to look through:
+  #
+  #     recording = TaskProgress::RecordingOutput.new
+  #     Importer.new(progress: TaskProgress.new("Import", outputs: [recording])).run
+  #     assert_includes recording.said, "Skipped 3 rows"
+  #     assert_equal 1, recording.events(:fail).length
+  #
+  # Safe to share between threads
+  class RecordingOutput
+    # @rbs @batches: Array[Batch]
+    # @rbs @lock: Thread::Mutex
+
+    # Nothing recorded yet
+    #: () -> void
+    def initialize
+      @batches = []
+      @lock = Mutex.new
+    end
+
+    # Records the batch
+    #: (Batch, Log) -> void
+    def call(batch, _log) = @lock.synchronize { @batches << batch }
+
+    # Every batch so far, in order
+    #: () -> Array[Batch]
+    def batches = @lock.synchronize { @batches.dup }
+
+    # Every event so far, in order - only those with the op `op`, if given
+    #: (?Symbol?) -> Array[Event]
+    def events(op = nil) = batches.flat_map(&:events).select { op.nil? || it.op == op }
+
+    # What the tree was told to `puts`, in order
+    #: () -> Array[String]
+    def said = events(:say).map { it.value.to_s }
+
+    # The tree the batches make, as an outline - see TaskProgress#outline
+    #: () -> Hash[Symbol, untyped]?
+    def outline
+      state = State.new
+      batches.each do |batch|
+        state = State.new if batch.reset?
+        batch.events.each { state.apply(it, at: batch.at) }
+      end
+      TaskProgress.outline(state, []) unless state.empty?
+    end
+  end
+end

data/lib/task_progress/redis_stream_store.rb

@@ -0,0 +1,131 @@
+# frozen_string_literal: true
+# rbs_inline: enabled
+
+class TaskProgress
+  # A stream store in Redis: each tree a Redis stream of its own (XADD, XREAD, XTRIM), its values in plain keys. It
+  # takes a redis-client connection or pool (`RedisClient.config(url:).new_pool`), or a redis-rb
+  # `Redis` - neither is a dependency of this gem, and nothing here loads one.
+  #
+  #     TaskProgress.stream_store = TaskProgress::RedisStreamStore.new(RedisClient.config(url: ENV["REDIS_URL"]).new_pool)
+  #
+  # Needs Redis 6.2 or later, for trimming by id. A reader that waits holds its connection for
+  # as long as it waits - a pool sized for the readers that follow at once keeps the rest moving.
+  # A redis-rb `Redis` is a single connection, so everything else on it waits too: give one that
+  # follows trees a client of its own.
+  #
+  # A write with a `guard:` - a key and the value it must hold - runs as a Lua script, so the
+  # check and the write are one step, see MemoryStreamStore.
+  #
+  # Every tree is under its own key already, so two users' trees never meet. A `namespace` puts
+  # all the keys under it as well - `tenant-42:task_progress:{...}` - for apps, or tenants, that
+  # share a Redis and must stay apart even so
+  class RedisStreamStore
+    # The guarded writes. KEYS[1] is the guard, ARGV[1] the value it must hold
+    GUARDED = {
+      set: "if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end " \
+        "if ARGV[3] == '' then redis.call('SET', KEYS[2], ARGV[2]) else redis.call('SET', KEYS[2], ARGV[2], 'PX', ARGV[3]) end return 1",
+      trim: "if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end redis.call('XTRIM', KEYS[2], 'MINID', ARGV[2]) return 1",
+      delete: "if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end for i = 2, #KEYS do redis.call('DEL', KEYS[i]) end return 1",
+      expire: "if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end redis.call('PEXPIRE', KEYS[2], ARGV[2]) return 1"
+    }.freeze #: Hash[Symbol, String]
+
+    # @rbs @redis: untyped
+    # @rbs @namespace: String?
+
+    # For a redis-client connection or pool, or a redis-rb client
+    #: (untyped, ?namespace: String?) -> void
+    def initialize(redis, namespace: nil)
+      @redis = redis
+      @namespace = namespace
+    end
+
+    # Adds `data` at the end of the stream at `key`, and returns its id
+    #: (String, String) -> String
+    def append(key, data) = @redis.call("XADD", ns(key), "*", "d", data).to_s
+
+    # The entries after the id `after` - all of them when it is nil - as `[id, data]` pairs. With
+    # a `timeout`, waits up to that many seconds for one to come when there are none yet
+    #: (String, ?after: String?, ?timeout: (Integer | Float)?) -> Array[[String, String]]
+    def read(key, after: nil, timeout: nil)
+      key = ns(key)
+      after ||= "0-0"
+      ms = (timeout * 1000).ceil if timeout&.positive?
+      reply = if ms.nil?
+        @redis.call("XREAD", "STREAMS", key, after)
+      elsif @redis.respond_to?(:blocking_call)
+        @redis.blocking_call(timeout.to_f + 1, "XREAD", "BLOCK", ms, "STREAMS", key, after)
+      else
+        # redis-rb knows to stretch its read timeout around a blocking command, but only in its own methods
+        @redis.xread(key, after, block: ms)
+      end
+      entries(reply, key)
+    end
+
+    # The value at `key`, or nil
+    #: (String) -> String?
+    def get(key) = @redis.call("GET", ns(key))
+
+    # Keeps `value` at `key`, for `expires_in` seconds if given
+    #: (String, String, ?expires_in: (Integer | Float)?, ?guard: [String, String]?) -> bool
+    def set(key, value, expires_in: nil, guard: nil)
+      key = ns(key)
+      ms = (expires_in * 1000).ceil if expires_in
+      return guarded(:set, guard, [key], [value, ms.to_s]) if guard
+      ms ? @redis.call("SET", key, value, "PX", ms) : @redis.call("SET", key, value)
+      true
+    end
+
+    # Forgets the streams and values at these keys
+    #: (*String, ?guard: [String, String]?) -> bool
+    def delete(*keys, guard: nil)
+      keys = keys.map { ns(it) }
+      return true if keys.empty?
+      return guarded(:delete, guard, keys, []) if guard
+      @redis.call("DEL", *keys)
+      true
+    end
+
+    # Drops the entries before the id `before`
+    #: (String, before: String, ?guard: [String, String]?) -> bool
+    def trim(key, before:, guard: nil)
+      key = ns(key)
+      return guarded(:trim, guard, [key], [before]) if guard
+      @redis.call("XTRIM", key, "MINID", before)
+      true
+    end
+
+    # Forgets what is at `key` - a stream or a value - `seconds` from now
+    #: (String, Integer | Float, ?guard: [String, String]?) -> bool
+    def expire(key, seconds, guard: nil)
+      key = ns(key)
+      ms = (seconds * 1000).ceil
+      return guarded(:expire, guard, [key], [ms.to_s]) if guard
+      @redis.call("PEXPIRE", key, ms)
+      true
+    end
+
+    private
+
+    #: (Symbol, [String, String], Array[String], Array[String]) -> bool
+    def guarded(script, guard, keys, args)
+      @redis.call("EVAL", GUARDED.fetch(script), keys.length + 1, ns(guard[0]), *keys, guard[1], *args).to_i == 1
+    end
+
+    #: (String) -> String
+    def ns(key) = @namespace ? "#{@namespace}:#{key}" : key
+
+    # XREAD answers in one of three shapes - RESP3 a Hash, RESP2 nested Arrays, redis-rb a Hash
+    # of Hashes - and nil when nothing came
+    #: (untyped, String) -> Array[[String, String]]
+    def entries(reply, key)
+      found = case reply
+      when Hash then reply[key]
+      when Array then reply.find { it[0] == key }&.last
+      end
+      (found || []).filter_map do |id, fields|
+        data = fields.is_a?(Hash) ? fields["d"] : fields.each_slice(2).find { it[0] == "d" }&.last
+        [id.to_s, data.to_s] if data
+      end
+    end
+  end
+end

data/lib/task_progress/squasher.rb

@@ -0,0 +1,59 @@
+# frozen_string_literal: true
+# rbs_inline: enabled
+
+class TaskProgress
+  # Squashes runs of batches into one. It mirrors the tree they build, so for a run it can
+  # send the prunes it saw and then a put for every node the run touched - which lands any
+  # reader that had the tree up to the run's `prev` exactly where the run left it, times
+  # and all, without replaying the steps in between. What was said is not in the tree, so
+  # it goes along after the puts
+  class Squasher # :nodoc:
+    # @rbs @state: State
+
+    # Starts from an empty tree
+    #: () -> void
+    def initialize = @state = State.new
+
+    # One batch that lands a reader where the run of batches would have
+    #: (Array[Batch]) -> Batch
+    def squash(batches)
+      prunes = [] #: Array[Event]
+      says = [] #: Array[Event]
+      touched = {} #: Hash[path, bool]
+      first = batches.first or raise ArgumentError, "Nothing to squash"
+      prev = first.prev
+      batches.each do |batch|
+        if batch.reset?
+          @state = State.new
+          prunes.clear
+          says.clear
+          touched.clear
+          prev = 0
+        end
+        journal = Journal.new
+        begin
+          batch.events.each { @state.apply(it, at: batch.at, journal:) }
+        rescue ArgumentError
+          # A batch that does not fit the tree it knows - it joined late, say - leaves that tree as
+          # it was. One alone goes on as it is, for whoever reads it to judge
+          @state.rollback(journal)
+          raise if batches.length > 1
+          return batch
+        end
+        batch.events.each do |event|
+          prunes << event if event.op == :prune
+          next says << event if event.op == :say
+          # set, add, note, done and fail start every node above theirs
+          (0..event.path.length).each { touched[event.path.take(it)] = true }
+        end
+      end
+      return first if batches.length == 1
+      last = batches.last or raise ArgumentError, "Nothing to squash"
+      puts = @state.each.filter_map { |path, node| Event.new(path:, op: :put, value: node.to_h) if touched[path] }
+      # The lines go after the puts, which bring their tasks back - but not for a task pruned
+      # since, which nothing after the prune would know
+      said = says.select { @state.key?(it.path) }
+      Batch.new(log: last.log, prev:, seq: last.seq, at: last.at, events: [*prunes, *puts, *said].freeze)
+    end
+  end
+end

data/lib/task_progress/sse.rb

@@ -0,0 +1,114 @@
+# frozen_string_literal: true
+# rbs_inline: enabled
+
+class TaskProgress
+  # A Rack response body streaming a tree as Server-Sent Events: first the tree so far, squashed
+  # into one batch with `prev` 0, then every batch as it commits - each in a `batch` event whose
+  # `id:` is its seq. A browser that reconnects starts over from a fresh squash, which is always
+  # right and costs a few puts. The stream ends once the root is finished. Given a subtask, it
+  # streams only that task's subtree, with the subtask as the root.
+  #
+  #     # Rack
+  #     TaskProgress::SSE.response(task)
+  #     # Rails
+  #     headers.merge!(TaskProgress::SSE::HEADERS)
+  #     self.response_body = TaskProgress::SSE.new(task)
+  #
+  # Given a key instead of a task - a tree handed off with TaskProgress.handoff - it follows that
+  # tree from the stream for as long as it streams, and stops following when the stream ends or
+  # the server closes it. `response` answers 404 for a key with nothing at it
+  class SSE
+    # What an SSE response is sent with
+    HEADERS = {"content-type" => "text/event-stream", "cache-control" => "no-cache", "x-accel-buffering" => "no"}.freeze #: Hash[String, String]
+
+    # @rbs @source: TaskProgress | Log
+    # @rbs @heartbeat: Integer | Float
+    # @rbs @reconnect: Integer | Float
+    # @rbs @close: bool
+
+    # A whole Rack response: status, headers and this body - or a 404, for a key with nothing at it
+    #: (TaskProgress | Log | String, **untyped) -> [Integer, Hash[String, String], SSE | Array[String]]
+    def self.response(source, **)
+      [200, HEADERS.dup, new(source, **)]
+    rescue KeyError
+      [404, {"content-type" => "text/plain"}, ["Nothing to follow"]]
+    end
+
+    # One batch as an SSE event
+    #: (Batch) -> String
+    def self.frame(batch) = "id: #{batch.seq}\nevent: batch\ndata: #{Wire.dump(batch)}\n\n"
+
+    # Streams the task - its subtree only, for a subtask - or the tree handed off at a key, from
+    # `stream`. It pings every `heartbeat` seconds of quiet, and tells the browser to reconnect
+    # after `reconnect` seconds
+    #: (TaskProgress | Log | String, ?stream_store: untyped, ?heartbeat: Integer | Float, ?reconnect: Integer | Float) -> void
+    def initialize(source, stream_store: TaskProgress.stream_store, heartbeat: 15, reconnect: 2)
+      @close = source.is_a?(String)
+      @source = source.is_a?(String) ? TaskProgress.follow(source, stream_store:) : source
+      @heartbeat = heartbeat
+      @reconnect = reconnect
+    end
+
+    # Stops following a tree it followed by its key - Rack calls it when the response is over
+    #: () -> void
+    def close
+      log.close if @close
+    end
+
+    # Yields the stream a chunk at a time, until the task finishes - as a Rack body does
+    #: () { (String) -> void } -> void
+    def each
+      queue = Thread::Queue.new
+      behind = false
+      # A client that stops reading would have every batch pile up here. Past BACKLOG it gets
+      # one squash of the tree instead, as of when it reads on
+      output = @source.add_output(replay: true) do |batch, _log|
+        next queue << batch if queue.size < BACKLOG
+        behind = true
+        queue.clear
+      end
+      yield "retry: #{(@reconnect * 1000).round}\n\n"
+      # A task pruned away before anyone asked: a tree with nothing in it, and the end
+      return yield(SSE.frame(Batch.new(log: log.id, prev: 0, seq: [log.seq, 1].max, at: 0.0, events: []))) unless present?
+      loop do
+        batch = queue.pop(timeout: @heartbeat)
+        if behind
+          behind = false
+          batch = log.squashed(path)
+        end
+        next yield(": ping\n\n") if batch.nil?
+        yield SSE.frame(batch)
+        break if finishes_root?(batch)
+      end
+    ensure
+      @source.remove_output(output) if output
+      close
+    end
+
+    # How many batches wait for a slow client before it gets a squash instead
+    BACKLOG = 1000 #: Integer
+
+    private
+
+    #: () -> Log
+    def log = (source = @source).is_a?(Log) ? source : source.log
+
+    #: () -> path
+    def path = (source = @source).is_a?(Log) ? [] : source.path
+
+    #: () -> bool
+    def present? = log.snapshot.key?(path)
+
+    # Judged by what was sent, never by the snapshot: that runs ahead of delivery, and ending on
+    # it could cut the stream before the batch that finished the root went out. An empty reset
+    # is a branch whose task was pruned, which will not come back either
+    #: (Batch) -> bool
+    def finishes_root?(batch)
+      return true if batch.reset? && batch.events.empty?
+      batch.events.any? do |event|
+        event.path.empty? && (%i[done fail cancel].include?(event.op) ||
+          (event.op == :put && %i[done failed canceled].include?(event.value[:status])))
+      end
+    end
+  end
+end