wsdirector-core 0.0.1 → 1.0.0

data/lib/wsdirector/protocols.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+require "wsdirector/protocols/base"
+require "wsdirector/protocols/action_cable"
+require "wsdirector/protocols/phoenix"
+
+module WSDirector
+  ID2CLASS = {
+    "base" => "Base",
+    "action_cable" => "ActionCable",
+    "phoenix" => "Phoenix"
+  }.freeze
+
+  module Protocols # :nodoc:
+    # Raised when received not expected message
+    class UnmatchedExpectationError < WSDirector::Error; end
+
+    # Raised when received message is unexpected
+    class UnexpectedMessageError < WSDirector::Error; end
+
+    # Raised when nothing has been received
+    class NoMessageError < WSDirector::Error; end
+
+    class << self
+      def get(id)
+        class_name = if ID2CLASS.key?(id)
+          ID2CLASS.fetch(id)
+        else
+          id
+        end
+
+        const_get(class_name)
+      rescue NameError
+        raise Error, "Unknown protocol: #{id}"
+      end
+    end
+  end
+end

data/lib/wsdirector/result.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+module WSDirector
+  # Handle results from all clients from the group
+  class Result
+    attr_reader :group, :errors
+
+    def initialize(group)
+      @group = group
+      @errors = Concurrent::Array.new
+
+      @all = Concurrent::AtomicFixnum.new(0)
+      @failures = Concurrent::AtomicFixnum.new(0)
+
+      @sampling_mutex = Mutex.new
+      @sampling_counter = Hash.new { |h, k| h[k] = 0 }
+    end
+
+    # Called when client successfully finished it's work
+    def succeed
+      all.increment
+    end
+
+    # Called when client failed
+    def failed(error_message)
+      errors << error_message
+      all.increment
+      failures.increment
+    end
+
+    def success?
+      failures.value.zero?
+    end
+
+    def total_count
+      all.value
+    end
+
+    def failures_count
+      failures.value
+    end
+
+    def track_sample(id, max)
+      sampling_mutex.synchronize do
+        return false if sampling_counter[id] >= max
+
+        sampling_counter[id] += 1
+        true
+      end
+    end
+
+    private
+
+    attr_reader :all, :success, :failures, :sampling_counter, :sampling_mutex
+  end
+end

data/lib/wsdirector/results_holder.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+module WSDirector
+  # Holds all results for all groups of clients
+  class ResultsHolder
+    def initialize
+      @groups = Concurrent::Map.new
+    end
+
+    def success?
+      @groups.values.all?(&:success?)
+    end
+
+    def groups
+      @groups.values
+    end
+
+    def print_summary(printer: $stdout, colorize: false)
+      single_group = groups.size == 1
+
+      @groups.each do |group, result|
+        color = result.success? ? :green : :red
+        prefix = single_group ? "" : "Group #{group}: "
+
+        msg = "#{prefix}#{result.total_count} clients, #{result.failures_count} failures\n"
+        msg = msg.colorize(color) if colorize
+
+        printer.puts(msg)
+
+        unless result.success?
+          print_errors(result.errors, printer: printer, colorize: colorize)
+          printer.puts "\n"
+        end
+      end
+    end
+
+    def <<(result)
+      @groups[result.group] = result
+    end
+
+    private
+
+    def print_errors(errors, printer:, colorize:)
+      printer.puts "\n"
+
+      errors.each.with_index do |error, i|
+        msg = "#{i + 1}) #{error}\n"
+        msg = msg.colorize(:red) if colorize
+
+        printer.puts msg
+      end
+    end
+  end
+end

data/lib/wsdirector/runner.rb

@@ -0,0 +1,60 @@
+# frozen_string_literal: true
+
+require "wsdirector/scenario_reader"
+require "wsdirector/clients_holder"
+require "wsdirector/results_holder"
+require "wsdirector/result"
+require "wsdirector/task"
+require "wsdirector/ext/deep_dup"
+
+module WSDirector
+  def self.run(scenario, scale: 1, connection_options: {}, locals: {}, **options)
+    scenario = ScenarioReader.parse(scenario, scale:, connection_options:, locals:)
+    Runner.new(scenario, scale:, **options).execute
+  end
+
+  # Initiates all clients as a separate tasks (=threads)
+  class Runner
+    using WSDirector::Ext::DeepDup
+
+    def initialize(scenario, url:, scale: 1, sync_timeout: 5, logger: nil, colorize: false)
+      @scenario = scenario
+      @url = url
+      @scale = scale
+      @logger = logger
+      @colorize = colorize
+      @total_count = scenario["total"]
+      @global_holder = ClientsHolder.new(total_count, sync_timeout:)
+      @results_holder = ResultsHolder.new
+    end
+
+    def execute
+      Thread.abort_on_exception = true
+
+      num = 0
+
+      tasks = scenario["clients"].flat_map do |client_config|
+        name = client_config.fetch("name")
+        result = Result.new(name)
+        results_holder << result
+
+        Array.new(client_config.fetch("multiplier")) do
+          num += 1
+          id = "#{name}_#{num}"
+          Thread.new do
+            Task.new(client_config.deep_dup, id:, colorize:, global_holder:, result:, scale:, logger:)
+              .run(url)
+          end
+        end
+      end
+
+      tasks.each(&:join)
+
+      results_holder
+    end
+
+    private
+
+    attr_reader :scenario, :url, :scale, :total_count, :global_holder, :results_holder, :logger, :colorize
+  end
+end

data/lib/wsdirector/scenario_reader.rb

@@ -0,0 +1,172 @@
+# frozen_string_literal: true
+
+require "erb"
+require "json"
+require "wsdirector/ext/deep_dup"
+
+module WSDirector
+  # Read and parse different scenarios
+  class ScenarioReader
+    using WSDirector::Ext::DeepDup
+
+    include WSDirector::Utils
+
+    class << self
+      def parse(scenario, **options)
+        new(**options).parse(scenario)
+      end
+    end
+
+    attr_reader :default_connections_options, :locals
+
+    def initialize(scale: 1, connection_options: {}, locals: {})
+      @scale = scale
+      @default_connections_options = connection_options
+      @locals = locals
+    end
+
+    def parse(scenario)
+      contents =
+        if scenario.is_a?(String)
+          if File.file?(scenario)
+            parse_file(scenario)
+          else
+            parse_from_str(scenario).then do
+              next _1 if _1.is_a?(Array)
+              [_1]
+            end
+          end.flatten
+        else
+          scenario
+        end
+
+      contents.map! do |item|
+        item.is_a?(String) ? {item => {}} : item
+      end
+
+      if contents.first&.key?("client")
+        contents = transform_with_loop(contents, multiple: true)
+        parse_multiple_scenarios(contents)
+      else
+        contents = transform_with_loop(contents)
+        {"total" => 1, "clients" => [parse_simple_scenario(contents)]}
+      end
+    end
+
+    private
+
+    JSON_FILE_FORMAT = /.+.(json)\z/
+    private_constant :JSON_FILE_FORMAT
+
+    def parse_from_str(contents)
+      JSON.parse(contents)
+    rescue JSON::ParserError
+      parse_yaml(contents)
+    end
+
+    def parse_file(file)
+      if file.match?(JSON_FILE_FORMAT)
+        JSON.parse(File.read(file))
+      else
+        parse_yaml(file)
+      end
+    end
+
+    def parse_yaml(path)
+      contents = File.file?(path) ? File.read(path) : path
+
+      if defined?(ERB)
+        contents = ERB.new(contents).result(erb_context)
+      end
+
+      ::YAML.load(contents, aliases: true, permitted_classes: [Date, Time, Regexp]) || {}
+    rescue ArgumentError
+      ::YAML.load(contents) || {}
+    end
+
+    def handle_steps(steps)
+      steps.flat_map.with_index do |step, id|
+        if step.is_a?(Hash)
+          type, data = step.to_a.first
+
+          data["sample"] = [1, parse_multiplier(data["sample"])].max if data["sample"]
+
+          multiplier = parse_multiplier(data.delete("multiplier") || "1")
+          Array.new(multiplier) { {"type" => type, "id" => id}.merge(data) }
+        else
+          {"type" => step, "id" => id}
+        end
+      end
+    end
+
+    def parse_simple_scenario(
+      steps,
+      multiplier: 1, name: "default", ignore: nil, protocol: "base",
+      connection_options: {}
+    )
+      {
+        "multiplier" => multiplier,
+        "steps" => handle_steps(steps),
+        "name" => name,
+        "ignore" => ignore,
+        "protocol" => protocol,
+        "connection_options" => default_connections_options.merge(connection_options)
+      }
+    end
+
+    def parse_multiple_scenarios(definitions)
+      total_count = 0
+      clients = definitions.map.with_index do |client, i|
+        _, client = client.to_a.first
+        multiplier = parse_multiplier(client.delete("multiplier") || "1")
+        name = client.delete("name") || (i + 1).to_s
+        connection_options = client.delete("connection_options") || {}
+        total_count += multiplier
+        ignore = parse_ignore(client.fetch("ignore", nil))
+        protocol = client.fetch("protocol", "base")
+
+        parse_simple_scenario(
+          client.fetch("actions", []),
+          multiplier:,
+          name:,
+          ignore:,
+          protocol:,
+          connection_options:
+        )
+      end
+      {"total" => total_count, "clients" => clients}
+    end
+
+    def transform_with_loop(contents, multiple: false)
+      contents.flat_map do |content|
+        loop_data = content.dig("client", "loop") || content.dig("loop")
+        next content unless loop_data
+
+        loop_multiplier = parse_multiplier(loop_data["multiplier"] || "1")
+
+        if multiple
+          content["client"]["actions"] = (loop_data["actions"] * loop_multiplier).map(&:deep_dup)
+          content
+        else
+          loop_data["actions"] * loop_multiplier
+        end
+      end
+    end
+
+    def parse_ignore(str)
+      return unless str
+
+      Array(str)
+    end
+
+    def erb_context
+      binding.then do |b|
+        locals.each do
+          b.local_variable_set(_1, _2)
+        end
+
+        b
+      end
+    end
+  end
+end

data/lib/wsdirector/snapshot.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module WSDirector
+  # Collect ws frame and dump them into a YAML file (or JSON)
+  class Snapshot
+    def initialize
+      @steps = []
+      @last_timestamp = nil
+    end
+
+    def <<(frame)
+      record_gap!
+      steps << {"send" => {"data" => frame}}
+    end
+
+    def to_yml
+      ::YAML.dump(steps)
+    end
+
+    def to_json
+      steps.to_json
+    end
+
+    private
+
+    attr_reader :steps, :last_timestamp
+
+    def record_gap!
+      prev_timestamp = last_timestamp
+      @last_timestamp = Time.now
+      return unless prev_timestamp
+
+      delay = last_timestamp - prev_timestamp
+      steps << {"sleep" => {"time" => delay}}
+    end
+  end
+end

data/lib/wsdirector/task.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+require "wsdirector/client"
+require "wsdirector/protocols"
+
+module WSDirector
+  # Single client operator
+  class Task
+    attr_reader :global_holder, :client
+
+    def initialize(config, id:, global_holder:, result:, scale:, logger:, colorize: false)
+      @id = id
+      @logger = logger
+      @ignore = config.fetch("ignore")
+      @steps = config.fetch("steps")
+      @connection_options = config.fetch("connection_options").transform_keys(&:to_sym)
+      @global_holder = global_holder
+      @result = result
+
+      protocol_class = Protocols.get(config.fetch("protocol", "base"))
+      @protocol = protocol_class.new(self, scale:, logger:, id:, color: color_for_id(id, colorize))
+    end
+
+    def run(url)
+      connect!(url, **@connection_options)
+
+      steps.each(&protocol)
+
+      result.succeed
+    rescue Error => e
+      result.failed(e.message)
+    end
+
+    def sampled?(step)
+      return true unless step["sample"]
+
+      id, max = step["id"], step["sample"]
+
+      result.track_sample(id, max)
+    end
+
+    private
+
+    attr_reader :steps, :result, :protocol, :id, :logger, :ignore
+
+    def connect!(url, **options)
+      protocol.init_client(url:, id:, ignore:, **options)
+    end
+
+    def color_for_id(id, colorize)
+      return unless colorize
+
+      String.colors[id.hash % String.colors.size]
+    end
+  end
+end

data/lib/wsdirector/utils.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module WSDirector
+  module Utils # :nodoc:
+    MULTIPLIER_FORMAT = /^[-+*\/\\\d ]+$/
+
+    attr_reader :scale
+
+    def parse_multiplier(str)
+      prepared = str.to_s.gsub(":scale", scale.to_s)
+      raise WSDirector::Error, "Unknown multiplier format: #{str}" unless
+        MULTIPLIER_FORMAT.match?(prepared)
+
+      eval(prepared) # rubocop:disable Security/Eval
+    end
+  end
+end

data/lib/wsdirector/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module WSDirector
+  VERSION = "1.0.0"
+end

data/lib/wsdirector-cli.rb

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+require "ruby-next"
+
+require "ruby-next/language/setup"
+RubyNext::Language.setup_gem_load_path(transpile: true)
+
+require "wsdirector/cli"

data/lib/wsdirector.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+require "concurrent"
+require "yaml"
+require "json"
+
+require "ruby-next"
+
+require "ruby-next/language/setup"
+RubyNext::Language.setup_gem_load_path(transpile: true)
+
+require "wsdirector/version"
+require "wsdirector/utils"
+
+# Command line tool for testing websocket servers using scenarios.
+module WSDirector
+  class Error < StandardError
+  end
+end
+
+require "wsdirector/runner"
+require "wsdirector/snapshot"