trifle.rb

require "connection_pool"

# Daylight-saving-safe timeline for trifle-stats reads.
#
# Trifle::Stats::Nocturnal#add walks a timeline by adding a fixed 86_400 seconds
# per day (604_800 per week). Tracking, on the other hand, floors the current
# time to local midnight. Once a read range crosses a daylight saving change,
# every generated bucket is an hour off the bucket the tracking wrote, no key
# matches, and the values come back as zeros - which is what the admin
# statistics showed for whole-year ranges.
#
# Advancing the wall clock in the configured zone keeps buckets on local
# midnight on both sides of the change. Re-check this patch when upgrading
# trifle-stats; as of 2.6.0 the fixed-seconds arithmetic is still there.
module TrifleStatsDstSafeTimeline
  def add(offset, unit)
    case unit
    when :day  then advance_wall_clock(days: offset)
    when :week then advance_wall_clock(days: offset * 7)
    else super
    end
  end

  private

    def advance_wall_clock(days:)
      raise ArgumentError, "Expected Time object, got #{time.class}" unless time.is_a?(Time)

      advanced = time.in_time_zone(config.time_zone).advance(days: days)

      Time.new(
        advanced.year, advanced.month, advanced.day,
        advanced.hour, advanced.min, advanced.sec, advanced.utc_offset
      )
    end
end

Trifle::Stats::Nocturnal.prepend(TrifleStatsDstSafeTimeline)

# Thread-safe, self-healing PostgreSQL connection for trifle-stats.
#
# trifle-stats' Postgres driver keeps a single PG::Connection and calls it
# directly on every operation. PG::Connection is not thread-safe, so under
# Puma's multi-threaded workers two threads sharing one connection corrupt the
# socket ("PQsocket() can't get socket descriptor"), and a connection that the
# server has dropped is never re-established.
#
# This wrapper exposes the methods the driver calls (each trifle operation makes
# a single top-level call to the client), but routes every call through a
# connection pool: each thread checks out its own connection, and a dead
# connection is reset and the call retried once before the error propagates.
class TrifleStatsConnection
  # Errors that mean the socket is gone and a reset+retry is worth attempting.
  RECOVERABLE = [ PG::ConnectionBad, PG::UnableToSend ].freeze

  def initialize(size:, timeout:, &connect)
    @pool = ConnectionPool.new(size: size, timeout: timeout, &connect)
  end

  def transaction(&block)
    with_connection { |conn| conn.transaction(&block) }
  end

  def exec_params(*args)
    with_connection { |conn| conn.exec_params(*args) }
  end

  def exec(*args)
    with_connection { |conn| conn.exec(*args) }
  end

  private

    def with_connection
      @pool.with do |conn|
        conn.reset unless conn.status == PG::CONNECTION_OK
        begin
          yield conn
        rescue *RECOVERABLE
          conn.reset
          yield conn
        end
      end
    end
end

Rails.application.config.after_initialize do
  next if ENV["SECRET_KEY_BASE_DUMMY"].present?

  db_config = ActiveRecord::Base.configurations.configs_for(env_name: Rails.env, name: "primary")

  Trifle::Stats.configure do |c|
    c.driver = if Rails.env.test?
      Trifle::Stats::Driver::Process.new
    else
      config_hash = db_config.configuration_hash
      pg_config = {
        host: config_hash[:host],
        port: config_hash[:port],
        dbname: config_hash[:database],
        user: config_hash[:username],
        password: config_hash[:password]
      }.compact

      # Size the pool to the worker's thread count so each thread can hold its
      # own connection without contention.
      pool_size = Integer(ENV.fetch("RAILS_MAX_THREADS", 3))
      client = TrifleStatsConnection.new(size: pool_size, timeout: 5) { PG.connect(pg_config) }

      Trifle::Stats::Driver::Postgres.new(client)
    end
    c.granularities = %w[10m 1h 1d]
    c.time_zone = "Europe/Bratislava"
    c.beginning_of_week = :monday
    c.buffer_enabled = false
  end
end