# To run this script you will need ruby install, then:
# 1. Ensure you have the 'pg' gem installed: gem install pg
# 2. Set environment variables for database connection if not using defaults:
#    PGHOST, PGPORT, PGDATABASE, PGUSER, PGPASSWORD
# 3. Execute the script: ruby workload_simulator_toast_lock_timeout.rb

require 'pg'
require 'thread'

# --- Configuration ---
DB_PARAMS = {
  host: ENV['PGHOST'] || 'localhost',
  port: ENV['PGPORT'] || 5432,
  dbname: ENV['PGDATABASE'] || 'postgres',
  user: ENV['PGUSER'] || 'shayon',
  password: ENV['PGPASSWORD']
}.freeze

TABLE_NAME = 'toast_target_table'.freeze
NUM_WORKER_THREADS = 15
WORK_DURATION_SECONDS = 7200
LOCK_TIMEOUT_MS = 500

def attempt_toast_access_with_timeout(conn, table_name, target_id, timeout_ms)
  begin
    conn.transaction do |txn_conn|
      txn_conn.exec("SET LOCAL lock_timeout = '#{timeout_ms}ms';")
      result = txn_conn.exec_params("SELECT length(payload) FROM #{table_name} WHERE id = $1", [target_id])
      puts "[#{Thread.current.object_id}] Accessed ID #{target_id}, result: #{result.to_a}"
    end
    true
  rescue PG::LockNotAvailable => e
    puts "[Worker #{Thread.current.object_id}] Lock timeout (#{timeout_ms}ms) accessing TOAST for ID #{target_id} on table #{table_name}."
    false
  rescue PG::QueryCanceled => e
    puts "[Worker #{Thread.current.object_id}] Query canceled (likely lock timeout: #{timeout_ms}ms) for ID #{target_id} on table #{table_name}."
    false
  rescue PG::ConnectionBad => e
    puts "[Worker #{Thread.current.object_id}] Connection bad: #{e.message.lines.first.strip}"
    raise
  rescue PG::Error => e
    puts "[Worker #{Thread.current.object_id}] Other PG::Error for ID #{target_id} on table #{table_name}: #{e.message.lines.first.strip}"
    false
  end
end

def worker_task(worker_id, table_name)
  conn = nil
  connection_retries = 0
  max_connection_retries = 5
  max_id_cache = nil
  last_max_id_check = Time.now - 60
  ops_this_connection = 0

  loop do
    begin
      conn = PG.connect(DB_PARAMS)
      connection_retries = 0
      ops_this_connection = 0
      start_time = Time.now

      while (Time.now - start_time) < WORK_DURATION_SECONDS
        if ops_this_connection % 500 == 0
            if !max_id_cache || (Time.now - last_max_id_check > 10)
                max_id_result = conn.query("SELECT MAX(id) FROM #{table_name}")
                if max_id_result && max_id_result.ntuples > 0 && max_id_result.getvalue(0,0)
                    current_max = max_id_result.getvalue(0,0).to_i
                    max_id_cache = current_max if current_max > 0
                    last_max_id_check = Time.now
                else
                    max_id_cache = nil
                end
            end
        end

        target_id = max_id_cache && max_id_cache > 0 ? rand(1..max_id_cache) : rand(1..100)

        attempt_toast_access_with_timeout(conn, table_name, target_id, LOCK_TIMEOUT_MS)

        ops_this_connection += 1
      end
      break
    rescue PG::ConnectionBad
      conn&.close
      connection_retries += 1
      if connection_retries > max_connection_retries
        puts "[Worker #{worker_id}] Max connection retries reached. Exiting."
        break
      end
      wait_time = 2**connection_retries
      puts "[Worker #{worker_id}] Connection lost. Retrying in #{wait_time}s..."
      sleep wait_time
    ensure
      conn&.close
    end
  end
end

# --- Main Execution ---
puts "Starting #{NUM_WORKER_THREADS} EXTREME TOAST ACCESS workload generator threads."
puts "Queries will have a lock timeout of #{LOCK_TIMEOUT_MS}ms."
puts "Targeting table: #{TABLE_NAME} on #{DB_PARAMS[:host]}:#{DB_PARAMS[:port]}/#{DB_PARAMS[:dbname]}"
puts "---"

threads = []
NUM_WORKER_THREADS.times do |i|
  threads << Thread.new do
    worker_task(i + 1, TABLE_NAME)
  end
end

threads.each(&:join)
puts "All workload threads finished."