discord_store 0.1.0

data/lib/discord_store/log.rb

@@ -0,0 +1,299 @@
+# frozen_string_literal: true
+
+require_relative "log/record"
+
+module DiscordStore
+  # An append-only log whose storage is a set of Discord channels.
+  #
+  # This is the layer everything else is built on, and it is append-only for a
+  # reason that is worth stating plainly, because it is the joke at the centre
+  # of this library.
+  #
+  # Deleting from Discord is expensive and gets worse with age: bulk deletion
+  # only covers messages under two weeks old, and past that it is one request
+  # per message against a rate limit that is already the binding constraint. So
+  # the cheap way to retire a record is to append a marker saying it is gone and
+  # let the reader skip it.
+  #
+  # That marker is a tombstone. Discord's own storage layer, Cassandra, did the
+  # same thing for the same reason — and the volume of tombstones a chat app
+  # generates is precisely what made Cassandra untenable for them and forced the
+  # migration to ScyllaDB. Building a store on Discord means writing tombstones
+  # into a database whose tombstones are stored as tombstones.
+  #
+  # Compaction is therefore not optional at scale; see {#compact}.
+  class Log
+    # @return [ChannelShard]
+    attr_reader :shard
+
+    # @return [Transport::REST]
+    attr_reader :rest
+
+    # @param rest [Transport::REST]
+    # @param config [Configuration]
+    # @param channel_ids [Array<String>, nil] defaults to config.log_channel_ids
+    # @param cipher [#seal, #open, nil]
+    def initialize(rest:, config:, channel_ids: nil, cipher: nil)
+      @rest = rest
+      @config = config
+      @shard = ChannelShard.new(channel_ids || config.log_channel_ids)
+      @cipher = cipher || Cipher.build(config)
+      @codec = Codec.new(cipher: @cipher, content_budget: config.content_budget)
+    end
+
+    # Appends one record.
+    #
+    # @param stream [String] logical topic and partition key
+    # @param data [Hash]
+    # @param nonce [String, nil] pass a stable value to make a retry idempotent
+    # @return [Log::Record] with +lsn+ assigned
+    def append(stream:, data:, nonce: nil)
+      append_all([Record.new(stream: stream, data: data, nonce: nonce)]).first
+    end
+
+    # Appends several records.
+    #
+    # Records are grouped by partition and packed, so N records cost far fewer
+    # than N requests. With +atomic+, every record must land in one message or
+    # the call raises: one message is the largest unit Discord commits
+    # all-or-nothing, and that is the only atomicity on offer here.
+    #
+    # @param records [Array<Log::Record>]
+    # @param atomic [Boolean]
+    # @return [Array<Log::Record>] the same records, with +lsn+ assigned
+    def append_all(records, atomic: false)
+      records = Array(records)
+      return [] if records.empty?
+
+      by_channel = records.group_by { |record| @shard.for(record.stream) }
+
+      if atomic && by_channel.size > 1
+        raise PayloadTooLargeError,
+              "an atomic batch spans #{by_channel.size} partitions (#{by_channel.keys.join(", ")}). " \
+              "Records that must commit together must share a partition; give them the same " \
+              "stream, or configure a single log channel."
+      end
+
+      by_channel.flat_map do |channel_id, channel_records|
+        write_to_channel(channel_id, channel_records, atomic: atomic)
+      end
+    end
+
+    # Buffers appends and writes them as one atomic batch on success.
+    #
+    #   log.transaction do |tx|
+    #     tx.append(stream: "orders", data: { id: 1 })
+    #     tx.append(stream: "orders", data: { id: 2 })
+    #   end
+    #
+    # Nothing is written if the block raises.
+    #
+    # @yieldparam buffer [Buffer]
+    # @return [Array<Log::Record>]
+    def transaction
+      buffer = Buffer.new
+      yield buffer
+      return [] if buffer.empty?
+
+      append_all(buffer.records, atomic: true)
+    end
+
+    # Reads records in partition order.
+    #
+    # @param stream [String, nil] a single stream, or nil for every partition
+    # @param after [String, nil] a cursor from {Record#cursor}; exclusive
+    # @param limit [Integer, nil]
+    # @yieldparam record [Log::Record]
+    # @return [Array<Log::Record>, void]
+    def each(stream: nil, after: nil, limit: nil, &block)
+      return enum_for(:each, stream: stream, after: after, limit: limit) unless block
+
+      channels = stream ? [@shard.for(stream)] : @shard.channel_ids
+      count = 0
+
+      channels.each do |channel_id|
+        read_channel(channel_id, after: after, stream: stream) do |record|
+          block.call(record)
+          count += 1
+          # Non-local on purpose: the limit applies to the whole walk across
+          # every partition, not to the current channel.
+          return if limit && count >= limit # rubocop:disable Lint/NonLocalExitFromIterator
+        end
+      end
+    end
+
+    # @return [Array<Log::Record>]
+    def read(stream: nil, after: nil, limit: nil)
+      each(stream: stream, after: after, limit: limit).to_a
+    end
+
+    # The cursor of the most recent record in a partition, or nil if empty.
+    #
+    # @param stream [String, nil]
+    # @return [String, nil]
+    def tip(stream: nil)
+      channel_id = stream ? @shard.for(stream) : @shard.channel_ids.first
+      message = @rest.list_messages(channel_id, limit: 1).first
+      return nil unless message && @codec.ours?(message)
+
+      records = decode(message)
+      records.last&.cursor
+    end
+
+    # Appends a tombstone for +target+.
+    #
+    # @param stream [String]
+    # @param target [String] the cursor or key being retired
+    # @return [Log::Record]
+    def tombstone(stream:, target:)
+      append_all([Record.tombstone(stream: stream, target: target)]).first
+    end
+
+    # Rewrites a message in place. Only valid for a message holding exactly one
+    # record; a packed message cannot be edited record-wise without rewriting
+    # its neighbours, and rewriting history in an append-only log is how replays
+    # start disagreeing with each other.
+    #
+    # @param record [Log::Record]
+    # @param data [Hash]
+    # @return [Log::Record]
+    def replace(record, data:)
+      raise ArgumentError, "record has no lsn" unless record.lsn
+
+      channel_id = record.channel_id || @shard.for(record.stream)
+      updated = Record.new(stream: record.stream, data: data, nonce: record.nonce,
+                           lsn: record.lsn, position: 0, channel_id: channel_id)
+
+      packed = @codec.pack([updated], aad: channel_id, atomic: true).first
+
+      unless packed[:files].empty?
+        raise PayloadTooLargeError, "replacement spills to an attachment; append a new record instead"
+      end
+
+      @rest.edit_message(channel_id, record.lsn, content: packed[:content])
+      updated
+    end
+
+    # Discards tombstoned records by physically deleting their messages.
+    #
+    # This is the expensive path, and the only one that reclaims anything. It is
+    # deliberately explicit rather than automatic: it costs one request per
+    # message once the two-week bulk-delete window has closed, which at the
+    # per-channel limit is roughly one message per second.
+    #
+    # @param stream [String]
+    # @param dry_run [Boolean] report what would be deleted without deleting
+    # @return [Hash] {scanned:, tombstoned:, deleted:, too_old:}
+    def compact(stream:, dry_run: false)
+      channel_id = @shard.for(stream)
+      targets = []
+      scanned = 0
+
+      read_channel(channel_id, stream: stream) do |record|
+        scanned += 1
+        targets << record.data["target"] if record.tombstone?
+      end
+
+      doomed = targets.filter_map { |cursor| cursor.to_s.split(":").first }.uniq
+
+      return { scanned: scanned, tombstoned: doomed.size, deleted: 0, too_old: 0, dry_run: true } if dry_run
+
+      too_old = @rest.bulk_delete_messages(channel_id, doomed)
+
+      if @config.delete_policy == :aggressive
+        too_old.each { |id| @rest.delete_message(channel_id, id) }
+        too_old = []
+      end
+
+      { scanned: scanned, tombstoned: doomed.size,
+        deleted: doomed.size - too_old.size, too_old: too_old.size, dry_run: false }
+    end
+
+    private
+
+    def write_to_channel(channel_id, records, atomic:)
+      packed = @codec.pack(records, aad: channel_id, atomic: atomic)
+
+      packed.flat_map do |message|
+        # One nonce per message, derived from the first record, so that a retry
+        # after a timeout is deduplicated by Discord rather than duplicated by us.
+        response = @rest.create_message(
+          channel_id,
+          content: message[:content],
+          nonce: message[:records].first.nonce,
+          files: message[:files]
+        )
+
+        message[:records].each_with_index.map do |record, position|
+          Record.new(stream: record.stream, data: record.data, nonce: record.nonce,
+                     lsn: response["id"], position: position, channel_id: channel_id)
+        end
+      end
+    end
+
+    def read_channel(channel_id, after: nil, stream: nil)
+      after_param, cursor_lsn, skip_position = parse_cursor(after)
+
+      @rest.each_message(channel_id, after: after_param) do |message|
+        next unless @codec.ours?(message)
+
+        decode(message).each do |record|
+          next if cursor_lsn && record.lsn == cursor_lsn && record.position <= skip_position
+          next if stream && record.stream != stream
+
+          yield record
+        end
+      end
+    end
+
+    def decode(message)
+      body = @codec.spilled?(message) ? fetch_spill(message) : nil
+      @codec.unpack(message, attachment_body: body, aad: message["channel_id"])
+    end
+
+    # Re-resolves the attachment URL from the message we are holding right now.
+    # A URL read from anywhere else is a URL that has probably expired.
+    def fetch_spill(message)
+      attachment = message["attachments"]&.first
+      unless attachment
+        raise CorruptRecordError,
+              "message #{message["id"]} claims a spill but has no attachment"
+      end
+
+      @rest.download(attachment["url"])
+    end
+
+    # Discord's +after+ is exclusive on message ID, so resuming inside a packed
+    # message means asking for the message itself back and then skipping the
+    # records already consumed.
+    #
+    # Returns three values, and the distinction between the first two matters:
+    # the bound sent to Discord is the cursor's message ID minus one, while the
+    # ID compared against each record is the cursor's own. Using the decremented
+    # value for both silently re-delivers the last record of every resume.
+    #
+    # @return [Array(String, String, Integer)] api bound, cursor lsn, skip position
+    def parse_cursor(cursor)
+      return [nil, nil, nil] if cursor.nil?
+
+      lsn, position = cursor.to_s.split(":")
+      [(lsn.to_i - 1).to_s, lsn, position.to_i]
+    end
+
+    # Collects records for {Log#transaction}.
+    class Buffer
+      attr_reader :records
+
+      def initialize = @records = []
+
+      def append(stream:, data:, nonce: nil)
+        record = Record.new(stream: stream, data: data, nonce: nonce)
+        @records << record
+        record
+      end
+
+      def empty? = @records.empty?
+      def size = @records.size
+    end
+  end
+end

data/lib/discord_store/railtie.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+require "rails/railtie"
+
+module DiscordStore
+  # Loads the rake tasks and the adapter into a Rails application.
+  #
+  #   # Gemfile
+  #   gem "discord_store", require: "discord_store/railtie"
+  class Railtie < ::Rails::Railtie
+    rake_tasks do
+      load File.expand_path("tasks.rake", __dir__)
+    end
+
+    initializer "discord_store.adapter" do
+      ActiveSupport.on_load(:active_record) do
+        require "active_record/connection_adapters/discord_adapter"
+      end
+    end
+
+    initializer "discord_store.active_storage" do
+      ActiveSupport.on_load(:active_storage_blob) do
+        require "active_storage/service/discord_service"
+      end
+    end
+  end
+end

data/lib/discord_store/replay.rb

@@ -0,0 +1,136 @@
+# frozen_string_literal: true
+
+module DiscordStore
+  # Rebuilds a local database from the Discord log.
+  #
+  # This is the operation that decides whether any of this is a real design or a
+  # stunt. If the channel is the source of truth, then a machine holding no
+  # local file must be able to reach the same state by reading it, and it must
+  # reach the *same* state every time. That is the whole contract:
+  #
+  #   rm storage/production.sqlite3 && rake discord:replay
+  #
+  # should leave the database exactly as it was, on any machine, from nothing
+  # but a bot token and a channel ID.
+  #
+  # Statements replay inside a local transaction per journalled transaction, so
+  # a batch that was atomic when written is atomic when replayed.
+  class Replay
+    # @param connection [ActiveRecord::ConnectionAdapters::AbstractAdapter]
+    # @param journal [Journal]
+    def initialize(connection:, journal:)
+      @connection = connection
+      @journal = journal
+    end
+
+    # @param from [String, nil] cursor to resume after; nil replays everything
+    # @param progress [#call, nil] called with (applied, cursor) periodically
+    # @return [Hash] {applied:, statements:, cursor:, skipped:}
+    def call(from: nil, progress: nil)
+      applied = 0
+      statements = 0
+      skipped = 0
+      cursor = from
+
+      # Replay is a rewrite of local state from the log, so it must not be
+      # journalled back into the log. Without this a replay would duplicate
+      # every transaction it read.
+      @journal.while_paused do
+        # Inside the pause, because reading the stored cursor creates the
+        # bookkeeping table, and that DDL must not become a log entry.
+        from ||= stored_cursor
+        cursor = from
+
+        @journal.each_transaction(from: from) do |entry, record_cursor|
+          begin
+            apply(entry["tx"])
+            applied += 1
+            statements += entry["tx"].size
+          rescue StandardError => e
+            skipped += 1
+            warn "discord_store: replay skipped #{record_cursor}: #{e.class}: #{e.message}"
+          end
+
+          cursor = record_cursor
+          store_cursor(cursor)
+          progress&.call(applied, cursor)
+        end
+      end
+
+      @journal.local_cursor = cursor
+      { applied: applied, statements: statements, cursor: cursor, skipped: skipped }
+    end
+
+    # How far behind the local database is.
+    #
+    # @return [Hash] {local:, remote:, behind:}
+    def status
+      local = stored_cursor
+      remote = @journal.tip
+      { local: local, remote: remote, behind: local != remote }
+    end
+
+    private
+
+    def apply(transaction)
+      @connection.transaction(requires_new: true) do
+        transaction.each do |statement|
+          sql = statement["sql"]
+          binds = Journal.deserialize_binds(statement["binds"])
+
+          if binds.empty?
+            @connection.execute(sql)
+          else
+            @connection.exec_query(sql, "Replay", type_casted_binds(sql, binds))
+          end
+        end
+      end
+    end
+
+    # Binds arrive as plain values, but exec_query wants attribute objects.
+    #
+    # Binary values need the binary type specifically: handed to the generic
+    # value type, SQLite's adapter tries to read them as UTF-8 and a byte like
+    # 0xFF ends the replay.
+    def type_casted_binds(_sql, binds)
+      binds.map do |value|
+        type = if value.is_a?(String) && (value.encoding == Encoding::BINARY || !value.valid_encoding?)
+                 ActiveModel::Type::Binary.new
+               else
+                 ActiveModel::Type::Value.new
+               end
+
+        ActiveRecord::Relation::QueryAttribute.new(nil, value, type)
+      end
+    end
+
+    def stored_cursor
+      ensure_cursor_table
+      row = @connection.select_one("SELECT cursor FROM #{Journal::CURSOR_TABLE} WHERE id = 1")
+      row && row["cursor"]
+    end
+
+    def store_cursor(cursor)
+      return if cursor.nil?
+
+      ensure_cursor_table
+      quoted = @connection.quote(cursor)
+      @connection.execute(
+        "INSERT INTO #{Journal::CURSOR_TABLE} (id, cursor) VALUES (1, #{quoted}) " \
+        "ON CONFLICT(id) DO UPDATE SET cursor = #{quoted}"
+      )
+    end
+
+    def ensure_cursor_table
+      return if @cursor_table_ready
+
+      @connection.execute(<<~SQL)
+        CREATE TABLE IF NOT EXISTS #{Journal::CURSOR_TABLE} (
+          id INTEGER PRIMARY KEY CHECK (id = 1),
+          cursor TEXT
+        )
+      SQL
+      @cursor_table_ready = true
+    end
+  end
+end

data/lib/discord_store/snowflake.rb

@@ -0,0 +1,109 @@
+# frozen_string_literal: true
+
+module DiscordStore
+  # Discord snowflake IDs, treated as what they actually are: a k-sortable
+  # primary key with a millisecond timestamp baked into the high bits.
+  #
+  #   63                                    22     17     12            0
+  #   +--------------------------------------+------+------+------------+
+  #   | milliseconds since 2015-01-01 (42b)  | wrkr | proc | incr (12b) |
+  #   +--------------------------------------+------+------+------------+
+  #
+  # Two properties matter for using Discord as a store:
+  #
+  # 1. IDs issued later sort after IDs issued earlier, so the natural message
+  #    order in a channel *is* insertion order, and no separate sequence column
+  #    is needed. The log sequence number is free.
+  #
+  # 2. Because the timestamp is recoverable, a wall-clock range maps onto an ID
+  #    range. Discord's +before+/+after+ pagination parameters take snowflakes,
+  #    so "every record written between 09:00 and 10:00" is a server-side range
+  #    scan rather than a client-side filter over the whole channel.
+  module Snowflake
+    # Discord's epoch: 2015-01-01T00:00:00Z, in milliseconds.
+    EPOCH_MS = 1_420_070_400_000
+
+    TIMESTAMP_SHIFT = 22
+    WORKER_SHIFT    = 17
+    PROCESS_SHIFT   = 12
+
+    WORKER_MASK    = 0x1F
+    PROCESS_MASK   = 0x1F
+    INCREMENT_MASK = 0xFFF
+
+    # The largest value the 42-bit timestamp field can hold.
+    MAX_TIMESTAMP_MS = (1 << 42) - 1
+
+    module_function
+
+    # Milliseconds since the Unix epoch at which +id+ was issued.
+    #
+    # @param id [Integer, String]
+    # @return [Integer]
+    def timestamp_ms(id)
+      (Integer(id) >> TIMESTAMP_SHIFT) + EPOCH_MS
+    end
+
+    # @param id [Integer, String]
+    # @return [Time] UTC time at which +id+ was issued
+    def at(id)
+      Time.at(timestamp_ms(id) / 1000.0).utc
+    end
+
+    # The worker/process/increment fields. Present for completeness and for
+    # debugging odd ordering; nothing in this library depends on them.
+    #
+    # @param id [Integer, String]
+    # @return [Hash]
+    def parts(id)
+      id = Integer(id)
+      {
+        timestamp_ms: timestamp_ms(id),
+        worker: (id >> WORKER_SHIFT) & WORKER_MASK,
+        process: (id >> PROCESS_SHIFT) & PROCESS_MASK,
+        increment: id & INCREMENT_MASK
+      }
+    end
+
+    # Build a synthetic snowflake for the given time, with all low bits zeroed.
+    #
+    # The result is not a real message ID and will never collide with one, but
+    # it sorts exactly where a real ID issued at that instant would sort. That
+    # makes it the correct bound for +before+/+after+ pagination when you want
+    # a time range rather than a message range.
+    #
+    # @param time [Time, Integer] a Time, or milliseconds since the Unix epoch
+    # @return [Integer]
+    # @raise [ArgumentError] if the time predates Discord or overflows 42 bits
+    def from_time(time)
+      ms = time.is_a?(Time) ? (time.to_f * 1000).floor : Integer(time)
+      offset = ms - EPOCH_MS
+
+      raise ArgumentError, "time predates the Discord epoch (2015-01-01)" if offset.negative?
+      raise ArgumentError, "time overflows the 42-bit snowflake timestamp" if offset > MAX_TIMESTAMP_MS
+
+      offset << TIMESTAMP_SHIFT
+    end
+
+    # The same instant as {.from_time}, but with every low bit set, so that it
+    # sorts after any real ID issued during that millisecond. Use this as an
+    # inclusive upper bound.
+    #
+    # @param time [Time, Integer]
+    # @return [Integer]
+    def from_time_inclusive(time)
+      from_time(time) | ((1 << TIMESTAMP_SHIFT) - 1)
+    end
+
+    # Whether +value+ could plausibly be a snowflake this library issued or read.
+    #
+    # @param value [Object]
+    # @return [Boolean]
+    def valid?(value)
+      id = Integer(value)
+      id.positive? && id < (1 << 64) && timestamp_ms(id) >= EPOCH_MS
+    rescue ArgumentError, TypeError
+      false
+    end
+  end
+end

data/lib/discord_store/tasks.rake

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+namespace :discord do
+  desc "Rebuild the local database from the Discord log"
+  task replay: :environment do
+    connection = ActiveRecord::Base.connection
+
+    unless connection.respond_to?(:replay!)
+      abort "The current connection is not a discord adapter (it is #{connection.adapter_name})."
+    end
+
+    started = Time.now
+    report = connection.replay!(from: ENV.fetch("FROM", nil))
+
+    puts "replayed #{report[:applied]} transactions (#{report[:statements]} statements) " \
+         "in #{(Time.now - started).round(1)}s"
+    puts "cursor now #{report[:cursor]}"
+    warn "#{report[:skipped]} transactions were skipped; see the log above" if report[:skipped].positive?
+  end
+
+  desc "Show how far the local database is behind the Discord log"
+  task status: :environment do
+    connection = ActiveRecord::Base.connection
+    abort "not a discord adapter" unless connection.respond_to?(:log_status)
+
+    status = connection.log_status
+    puts "journal mode:  #{status[:mode]}"
+    puts "local cursor:  #{status[:local_cursor] || "(never replayed)"}"
+    puts "remote cursor: #{status[:remote_cursor] || "(empty log)"}"
+    puts "buffered:      #{status[:pending]} transactions not yet written" if status[:pending].to_i.positive?
+
+    if status[:local_cursor] == status[:remote_cursor]
+      puts "up to date."
+    else
+      puts "behind. Run rake discord:replay."
+    end
+  end
+
+  desc "Delete tombstoned records from a stream (STREAM=name, DRY_RUN=1)"
+  task compact: :environment do
+    stream = ENV.fetch("STREAM", nil) or abort "STREAM=name is required"
+    dry_run = ENV["DRY_RUN"] == "1"
+
+    report = DiscordStore.client.log.compact(stream: stream, dry_run: dry_run)
+
+    puts "scanned #{report[:scanned]} records, #{report[:tombstoned]} tombstoned"
+    if dry_run
+      puts "(dry run; nothing deleted)"
+    else
+      puts "deleted #{report[:deleted]} messages"
+      if report[:too_old].positive?
+        puts "#{report[:too_old]} were older than the 14-day bulk-delete window and were left in " \
+             "place. Re-run with delete_policy: :aggressive to remove them one request at a time."
+      end
+    end
+  end
+
+  desc "Generate an encryption key"
+  task :key do
+    puts DiscordStore.generate_key
+  end
+
+  desc "Check the token, the application id, and the channels"
+  task doctor: :environment do
+    client = DiscordStore.client
+    user = client.verify!
+    puts "token ok: #{user["id"]} (#{user["username"] || "bot"})"
+
+    client.config.all_channel_ids.each do |channel_id|
+      client.rest.get_channel(channel_id)
+      puts "channel #{channel_id}: reachable"
+    rescue DiscordStore::APIError => e
+      puts "channel #{channel_id}: #{e.status} #{e.message}"
+    end
+
+    puts "attachment limit: #{client.limits.attachment_limit} bytes"
+    puts "chunk size:       #{client.limits.chunk_size} bytes"
+  rescue DiscordStore::Error => e
+    abort "#{e.class}: #{e.message}"
+  end
+
+  desc "Measure the real attachment ceiling by uploading test files (CHANNEL=id)"
+  task probe: :environment do
+    channel = ENV["CHANNEL"] || DiscordStore.client.config.blob_channel_ids.first
+    abort "CHANNEL=id is required" unless channel
+
+    puts "probing #{channel}; this uploads and deletes throwaway attachments..."
+    limit = DiscordStore.client.limits.probe!(channel_id: channel)
+    puts "largest accepted attachment: #{limit} bytes (#{(limit / 1024.0 / 1024).round(2)} MiB)"
+    puts "pin it: config.chunk_size = #{limit - DiscordStore::GuildLimits::REQUEST_OVERHEAD}"
+  end
+end