mahout 1.1.0

data/lib/mahout/config.rb

@@ -0,0 +1,322 @@
+# frozen_string_literal: true
+
+require "ipaddr"
+
+module Mahout
+  class Config
+    REQUIRED_KEYS = %w[server instance postgres storage hardening].freeze
+    VALID_PROFILES = %w[conservative balanced aggressive].freeze
+    SAFE_IDENTIFIER = /\A[a-zA-Z_][a-zA-Z0-9_]*\z/
+
+    attr_reader :data
+
+    def initialize(path: nil, data: nil, merge: nil, overrides: {})
+      @data = if data
+                data
+              elsif path
+                YAML.safe_load_file(path, permitted_classes: [Symbol])
+              else
+                raise ConfigError, "must provide either path or data"
+              end
+      deep_merge!(@data, merge) if merge
+      merge_overrides(overrides)
+      validate
+    end
+
+    def checksum
+      Digest::MD5.hexdigest(@data.to_json)
+    end
+
+    def set_secret(key, value)
+      @secrets[key] = value
+    end
+
+    def secret(key)
+      @secrets[key]
+    end
+
+    def host
+      dig("server", "host")
+    end
+
+    def user
+      dig("server", "user") || "root"
+    end
+
+    def key
+      dig("server", "key") || "~/.ssh/id_rsa"
+    end
+
+    def ssh_options
+      opts = dig("server", "ssh_options") || {}
+      opts.transform_keys(&:to_sym)
+    end
+
+    def instance_name
+      dig("instance", "name")
+    end
+
+    def environment
+      dig("instance", "environment") || "production"
+    end
+
+    def pg_version
+      dig("postgres", "version") || 17
+    end
+
+    def pg_port
+      dig("postgres", "port") || 5432
+    end
+
+    def pg_database
+      dig("postgres", "database")
+    end
+
+    def pg_username
+      dig("postgres", "username")
+    end
+
+    def profile
+      dig("profile") || "balanced"
+    end
+
+    def data_device
+      dig("storage", "data_device")
+    end
+
+    def data_mount
+      dig("storage", "data_mount")
+    end
+
+    def wal_device
+      dig("storage", "wal_device")
+    end
+
+    def wal_mount
+      dig("storage", "wal_mount")
+    end
+
+    def pgbouncer_port
+      dig("pgbouncer", "port") || 6432
+    end
+
+    def pgbouncer_pool_mode
+      dig("pgbouncer", "pool_mode") || "transaction"
+    end
+
+    def pgbackrest_enabled?
+      dig("pgbackrest", "enabled") != false && dig("pgbackrest") != nil
+    end
+
+    def pgbackrest_stanza
+      dig("pgbackrest", "stanza")
+    end
+
+    def pgbackrest_s3_bucket
+      dig("pgbackrest", "s3_bucket")
+    end
+
+    def pgbackrest_s3_endpoint
+      endpoint = dig("pgbackrest", "s3_endpoint") || "s3.amazonaws.com"
+      endpoint.sub(%r{\Ahttps?://}, "")
+    end
+
+    def pgbackrest_s3_region
+      dig("pgbackrest", "s3_region") || "us-east-1"
+    end
+
+    def pgbackrest_s3_uri_style
+      dig("pgbackrest", "s3_uri_style") || "host"
+    end
+
+    def pgbackrest_retention_full
+      dig("pgbackrest", "retention_full") || 4
+    end
+
+    def pgbackrest_process_max
+      dig("pgbackrest", "process_max") || 4
+    end
+
+    def pgbackrest_schedule_full
+      dig("pgbackrest", "schedule", "full") || "Sun *-*-* 02:00:00"
+    end
+
+    def pgbackrest_schedule_diff
+      dig("pgbackrest", "schedule", "diff") || "Mon..Sat *-*-* 02:00:00"
+    end
+
+    def pgbackrest_schedule_incr
+      dig("pgbackrest", "schedule", "incr") || "*-*-* *:30:00"
+    end
+
+    def datadog_enabled?
+      dig("datadog", "enabled") == true
+    end
+
+    def datadog_site
+      dig("datadog", "site") || "datadoghq.com"
+    end
+
+    def datadog_tags
+      dig("datadog", "tags") || []
+    end
+
+    def allowed_cidrs
+      dig("hardening", "allowed_cidrs") || []
+    end
+
+    def ssl_required?
+      dig("hardening", "ssl_required") != false
+    end
+
+    def ssh_hardening?
+      dig("hardening", "ssh_hardening") != false
+    end
+
+    def ssh_password_auth?
+      dig("hardening", "ssh_password_auth") == true
+    end
+
+    def ssh_max_auth_tries
+      dig("hardening", "ssh_max_auth_tries") || 3
+    end
+
+    def fail2ban?
+      dig("hardening", "fail2ban") != false
+    end
+
+    def fail2ban_maxretry
+      dig("hardening", "fail2ban_maxretry") || 5
+    end
+
+    def fail2ban_bantime
+      dig("hardening", "fail2ban_bantime") || 3600
+    end
+
+    def unattended_upgrades?
+      dig("hardening", "unattended_upgrades") != false
+    end
+
+    def ufw?
+      dig("hardening", "ufw") != false
+    end
+
+    def templates_path
+      path = dig("templates_path")
+      return nil unless path
+
+      File.expand_path(path)
+    end
+
+    def extensions
+      dig("extensions") || ["pg_stat_statements"]
+    end
+
+    def postgres_overrides
+      dig("postgres", "overrides") || {}
+    end
+
+    def tuner_profile
+      dig("tuner", "profile")
+    end
+
+    def pg_data_dir
+      "#{data_mount}/#{pg_version}/main"
+    end
+
+    def pg_wal_dir
+      return nil unless wal_mount
+
+      "#{wal_mount}/#{pg_version}/wal"
+    end
+
+    private
+
+    def deep_merge!(base, overlay)
+      overlay.each do |key, value|
+        if value.is_a?(Hash) && base[key].is_a?(Hash)
+          deep_merge!(base[key], value)
+        else
+          base[key] = value
+        end
+      end
+    end
+
+    def dig(*keys)
+      @data.dig(*keys)
+    end
+
+    def merge_overrides(overrides)
+      @secrets = {}
+      load_secrets_file
+      overrides.each do |key, value|
+        case key.to_s
+        when "host"
+          @data["server"] ||= {}
+          @data["server"]["host"] = value
+        when "user"
+          @data["server"] ||= {}
+          @data["server"]["user"] = value
+        when "key"
+          @data["server"] ||= {}
+          @data["server"]["key"] = value
+        end
+      end
+    end
+
+    def load_secrets_file
+      secrets_path = dig("secrets_file")
+      return unless secrets_path
+
+      path = File.expand_path(secrets_path)
+      return unless File.exist?(path)
+
+      secrets_data = YAML.safe_load_file(path, permitted_classes: [Symbol])
+      return unless secrets_data.is_a?(Hash)
+
+      secrets_data.each { |k, v| @secrets[k.to_s] = v.to_s }
+    end
+
+    def validate
+      REQUIRED_KEYS.each do |key|
+        raise ConfigError, "missing required config section: #{key}" unless @data[key]
+      end
+      raise ConfigError, "missing server.host" unless host
+      raise ConfigError, "missing postgres.database" unless pg_database
+      raise ConfigError, "missing postgres.username" unless pg_username
+      raise ConfigError, "missing storage.data_device" unless data_device
+      raise ConfigError, "missing storage.data_mount" unless data_mount
+      if pgbackrest_enabled?
+        raise ConfigError, "missing pgbackrest.stanza" unless pgbackrest_stanza
+        raise ConfigError, "missing pgbackrest.s3_bucket" unless pgbackrest_s3_bucket
+      end
+
+      validate_safe_identifier("postgres.database", pg_database)
+      validate_safe_identifier("postgres.username", pg_username)
+      validate_cidrs
+      validate_profile
+    end
+
+    def validate_safe_identifier(label, value)
+      return if value.nil?
+
+      unless value.match?(SAFE_IDENTIFIER)
+        raise ConfigError, "#{label} contains unsafe characters: #{value}. only alphanumeric and underscore allowed"
+      end
+    end
+
+    def validate_cidrs
+      allowed_cidrs.each do |cidr|
+        IPAddr.new(cidr)
+      rescue IPAddr::InvalidAddressError
+        raise ConfigError, "invalid CIDR: #{cidr}"
+      end
+    end
+
+    def validate_profile
+      return if VALID_PROFILES.include?(profile)
+
+      raise ConfigError, "unknown profile: #{profile}. must be one of: #{VALID_PROFILES.join(", ")}"
+    end
+  end
+end

data/lib/mahout/disk_benchmark.rb

@@ -0,0 +1,172 @@
+# frozen_string_literal: true
+
+module Mahout
+  class DiskBenchmark
+    BLOCK_SIZE = "1M"
+    COUNT = 1024
+    TEST_FILE = "mahout-diskbench"
+
+    def initialize(runner:, config:)
+      @runner = runner
+      @config = config
+    end
+
+    def call
+      data_mount = @config.data_mount
+      wal_mount = @config.wal_mount
+
+      ensure_fio
+
+      $stdout.puts("disk benchmark")
+
+      $stdout.puts("")
+      $stdout.puts("sequential throughput (1M blocks):")
+      run_throughput_tests("data (#{data_mount})", data_mount)
+
+      if wal_mount
+        $stdout.puts("")
+        run_throughput_tests("wal (#{wal_mount})", wal_mount)
+        $stdout.puts("")
+        run_cross_device(wal_mount, data_mount)
+      end
+
+      $stdout.puts("")
+      $stdout.puts("random iops (4k blocks):")
+      run_iops_tests("data (#{data_mount})", data_mount)
+
+      if wal_mount
+        $stdout.puts("")
+        run_iops_tests("wal (#{wal_mount})", wal_mount)
+      end
+
+      $stdout.puts("")
+      $stdout.puts("done")
+    end
+
+    private
+
+    def ensure_fio
+      result = @runner.run("which fio", allow_failure: true)
+      return if result.success? || @runner.dry_run?
+
+      @runner.run("apt-get install -y fio", allow_failure: true)
+    end
+
+    def run_throughput_tests(label, mount)
+      $stdout.puts("#{label}:")
+      run_write(mount)
+      clear_cache
+      run_read(mount)
+      cleanup(mount)
+    end
+
+    def run_iops_tests(label, mount)
+      $stdout.puts("#{label}:")
+      file = "#{mount}/#{TEST_FILE}-fio"
+
+      run_fio(label, mount, "randread", file)
+      run_fio(label, mount, "randwrite", file)
+      run_fio(label, mount, "randrw", file)
+
+      @runner.run("rm -f #{file}", allow_failure: true)
+    end
+
+    def run_fio(label, mount, mode, file)
+      result = @runner.run(
+        "fio --name=test --filename=#{file} --size=1G --bs=4k " \
+        "--rw=#{mode} --direct=1 --ioengine=libaio --iodepth=64 " \
+        "--numjobs=4 --group_reporting --runtime=10 --time_based " \
+        "--output-format=json 2>/dev/null",
+        allow_failure: true
+      )
+      print_fio_result("  #{mode}", result)
+    end
+
+    def print_fio_result(label, result)
+      return if @runner.dry_run?
+      return $stdout.puts("#{label}: could not run") unless result.success?
+
+      data = JSON.parse(result.stdout) rescue return
+      job = data.dig("jobs", 0)
+      return unless job
+
+      read_iops = job.dig("read", "iops")&.round
+      write_iops = job.dig("write", "iops")&.round
+      read_lat = job.dig("read", "lat_ns", "mean")
+      write_lat = job.dig("write", "lat_ns", "mean")
+
+      parts = []
+      if read_iops && read_iops > 0
+        lat_us = read_lat ? " (#{(read_lat / 1000.0).round(1)}us)" : ""
+        parts << "read #{read_iops} iops#{lat_us}"
+      end
+      if write_iops && write_iops > 0
+        lat_us = write_lat ? " (#{(write_lat / 1000.0).round(1)}us)" : ""
+        parts << "write #{write_iops} iops#{lat_us}"
+      end
+
+      $stdout.puts("#{label}: #{parts.join(", ")}")
+    end
+
+    def run_write(mount)
+      file = "#{mount}/#{TEST_FILE}"
+      result = @runner.run(
+        "dd if=/dev/zero of=#{file} bs=#{BLOCK_SIZE} count=#{COUNT} conv=fdatasync 2>&1",
+        allow_failure: true
+      )
+      print_dd_result("  write", result)
+    end
+
+    def run_read(mount)
+      file = "#{mount}/#{TEST_FILE}"
+      result = @runner.run(
+        "dd if=#{file} of=/dev/null bs=#{BLOCK_SIZE} count=#{COUNT} 2>&1",
+        allow_failure: true
+      )
+      print_dd_result("  read", result)
+    end
+
+    def run_cross_device(from_mount, to_mount)
+      from_file = "#{from_mount}/#{TEST_FILE}"
+      to_file = "#{to_mount}/#{TEST_FILE}-cross"
+
+      @runner.run(
+        "dd if=/dev/zero of=#{from_file} bs=#{BLOCK_SIZE} count=#{COUNT} conv=fdatasync 2>&1",
+        allow_failure: true
+      )
+      clear_cache
+
+      $stdout.puts("wal -> data (#{from_mount} -> #{to_mount}):")
+      result = @runner.run(
+        "dd if=#{from_file} of=#{to_file} bs=#{BLOCK_SIZE} count=#{COUNT} conv=fdatasync 2>&1",
+        allow_failure: true
+      )
+      print_dd_result("  copy", result)
+
+      @runner.run("rm -f #{from_file} #{to_file}", allow_failure: true)
+    end
+
+    def clear_cache
+      @runner.run("sync && echo 3 > /proc/sys/vm/drop_caches", allow_failure: true)
+    end
+
+    def cleanup(mount)
+      @runner.run("rm -f #{mount}/#{TEST_FILE}", allow_failure: true)
+    end
+
+    def print_dd_result(label, result)
+      return if @runner.dry_run?
+
+      output = result.stdout
+      speed_line = output.lines.find { |l| l.include?("/s") }
+      if speed_line
+        speed = speed_line.strip.split(",").last.strip
+        bytes = speed_line.strip.match(/(\d[\d.]*\s*[kMGT]?B)\s+copied/)
+        size = bytes ? bytes[1] : "#{COUNT * 1} MB"
+        $stdout.puts("#{label}: #{size}, #{speed}")
+      else
+        $stdout.puts("#{label}: #{output.lines.last&.strip}")
+      end
+    end
+  end
+end

data/lib/mahout/extension_registry.rb

@@ -0,0 +1,144 @@
+# frozen_string_literal: true
+
+module Mahout
+  class ExtensionRegistry
+    Extension = Data.define(
+      :name,
+      :package,
+      :preload,
+      :depends_on,
+      :apt_repo,
+      :create_options
+    )
+
+    EXTENSIONS = {
+      "pg_stat_statements" => Extension.new(
+        name: "pg_stat_statements",
+        package: nil,
+        preload: true,
+        depends_on: [],
+        apt_repo: nil,
+        create_options: nil
+      ),
+      "pg_trgm" => Extension.new(
+        name: "pg_trgm",
+        package: nil,
+        preload: false,
+        depends_on: [],
+        apt_repo: nil,
+        create_options: nil
+      ),
+      "btree_gist" => Extension.new(
+        name: "btree_gist",
+        package: nil,
+        preload: false,
+        depends_on: [],
+        apt_repo: nil,
+        create_options: nil
+      ),
+      "pgcrypto" => Extension.new(
+        name: "pgcrypto",
+        package: nil,
+        preload: false,
+        depends_on: [],
+        apt_repo: nil,
+        create_options: nil
+      ),
+      "postgis" => Extension.new(
+        name: "postgis",
+        package: "postgresql-%{version}-postgis-3",
+        preload: false,
+        depends_on: [],
+        apt_repo: nil,
+        create_options: nil
+      ),
+      "timescaledb" => Extension.new(
+        name: "timescaledb",
+        package: "timescaledb-2-postgresql-%{version}",
+        preload: true,
+        depends_on: [],
+        apt_repo: "deb https://packagecloud.io/timescale/timescaledb/ubuntu/ %{codename} main",
+        create_options: nil
+      ),
+      "pg_cron" => Extension.new(
+        name: "pg_cron",
+        package: "postgresql-%{version}-cron",
+        preload: true,
+        depends_on: [],
+        apt_repo: nil,
+        create_options: nil
+      ),
+      "pg_partman" => Extension.new(
+        name: "pg_partman",
+        package: "postgresql-%{version}-partman",
+        preload: true,
+        depends_on: [],
+        apt_repo: nil,
+        create_options: nil
+      ),
+      "hstore" => Extension.new(
+        name: "hstore",
+        package: nil,
+        preload: false,
+        depends_on: [],
+        apt_repo: nil,
+        create_options: nil
+      ),
+      "pg_repack" => Extension.new(
+        name: "pg_repack",
+        package: "postgresql-%{version}-repack",
+        preload: true,
+        depends_on: [],
+        apt_repo: nil,
+        create_options: nil
+      ),
+      "pgvector" => Extension.new(
+        name: "vector",
+        package: "postgresql-%{version}-pgvector",
+        preload: false,
+        depends_on: [],
+        apt_repo: nil,
+        create_options: nil
+      ),
+      "pg_prewarm" => Extension.new(
+        name: "pg_prewarm",
+        package: nil,
+        preload: true,
+        depends_on: [],
+        apt_repo: nil,
+        create_options: nil
+      )
+    }.freeze
+
+    class << self
+      def resolve(names, pg_version:)
+        resolved = []
+        names.each do |name|
+          ext = EXTENSIONS.fetch(name) do
+            raise UnknownExtension, "unknown extension: #{name}. known: #{EXTENSIONS.keys.join(", ")}"
+          end
+          ext.depends_on.each { |dep| resolved << EXTENSIONS.fetch(dep) }
+          resolved << ext
+        end
+        resolved.uniq(&:name)
+      end
+
+      def packages(extensions, pg_version:)
+        extensions
+          .filter_map { |ext| ext.package&.% ({ version: pg_version }) }
+          .uniq
+      end
+
+      def preload_libraries(extensions)
+        extensions.select(&:preload).map(&:name)
+      end
+
+      def apt_repos(extensions, pg_version:, codename:)
+        extensions
+          .filter_map(&:apt_repo)
+          .map { |repo| repo % { version: pg_version, codename: codename } }
+          .uniq
+      end
+    end
+  end
+end