container_broker 1.0.1

data/app/services/check_for_slot_removal.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+class CheckForSlotRemoval
+  attr_reader :slot
+
+  def initialize(slot:)
+    @slot = slot
+  end
+
+  def perform
+    return unless adjust_execution_type_slots_instance.decrement?
+
+    slot.destroy!
+    FriendlyNameSlots.new(node: slot.node.reload).perform
+
+    @removed = true
+  end
+
+  def removed?
+    @removed
+  end
+
+  private
+
+  def adjust_execution_type_slots_instance
+    AdjustExecutionTypeSlots.new(node: slot.node, execution_type: slot.execution_type)
+  end
+end

data/app/services/collect_load_metrics.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+class CollectLoadMetrics
+  def perform
+    send_tasks_count
+    send_slots_count
+    send_slots_usage_percent
+  end
+
+  private
+
+  def send_tasks_count
+    send_metrics("tasks_count", Task.where(:status.in => %w[waiting starting started retry failed]))
+  end
+
+  def send_slots_count
+    send_metrics("slots_count", Node.available.flat_map(&:slots))
+  end
+
+  def send_slots_usage_percent
+    Node
+      .available
+      .flat_map(&:slots)
+      .group_by(&:execution_type)
+      .each do |execution_type, slots|
+        Metrics.new("slots_usage_percent").count(execution_type: execution_type, percent: SlotsUsagePercentage.new(slots).perform)
+      end
+  end
+
+  def send_metrics(metric, items)
+    items
+      .group_by { |s| { execution_type: s.execution_type, status: s.status } }
+      .transform_values(&:count)
+      .each do |data, count|
+        Metrics.new(metric).count(
+          data.merge(amount: count)
+        )
+      end
+  end
+end

data/app/services/delete_node.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+class DeleteNode
+  class NodeWithRunningSlotsError < StandardError; end
+
+  attr_reader :node
+
+  def initialize(node:)
+    @node = node
+  end
+
+  def perform
+    if node.accept_new_tasks
+      was_accepting_new_tasks = true
+      NodeTaskAcceptance.new(node: node).reject!
+    end
+
+    if node.slots.working.any?
+      NodeTaskAcceptance.new(node: node).accept! if was_accepting_new_tasks
+      raise NodeWithRunningSlotsError
+    else
+      node.destroy!
+    end
+  end
+end

data/app/services/friendly_name_nodes.rb

@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+class FriendlyNameNodes
+  def perform
+    Node.order(runner_provider: :desc, hostname: :asc, id: :asc).each_with_index do |node, index|
+      node.update(name: "n#{format("%02d%s", (index + 1), node.runner_provider.first)}")
+      FriendlyNameSlots.new(node: node).perform
+    end
+  end
+end

data/app/services/friendly_name_slots.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+class FriendlyNameSlots
+  attr_reader :node
+
+  def initialize(node:)
+    @node = node
+  end
+
+  def perform
+    node.slots.each_with_index do |slot, index|
+      slot.update(name: "#{node.name}-s#{format("%02d", (index + 1))}-#{slot.execution_type}")
+    end
+  end
+end

data/app/services/kill_node_runners.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+class KillNodeRunners
+  attr_reader :node
+
+  def initialize(node:)
+    @node = node
+  end
+
+  def perform
+    node.slots.running.each do |slot|
+      node
+        .runner_service(:kill_slot_runner)
+        .perform(slot: slot)
+    end
+  end
+end

data/app/services/kill_task_container.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+class KillTaskContainer
+  attr_reader :task
+
+  class TaskNotRunningError < StandardError; end
+
+  delegate :node, to: :task
+
+  def initialize(task:)
+    @task = task
+  end
+
+  def perform
+    validate_task_status
+
+    task
+      .slot
+      .node
+      .runner_service(:kill_slot_runner)
+      .perform(slot: task.slot)
+  end
+
+  private
+
+  def validate_task_status
+    raise TaskNotRunningError, "#{task} is not running" unless task.started?
+  end
+end

data/app/services/kubernetes_client.rb

@@ -0,0 +1,136 @@
+# frozen_string_literal: true
+
+class KubernetesClient
+  class PodNotFoundError < StandardError; end
+  class NetworkError < StandardError; end
+  class LogsNotFoundError < StandardError; end
+
+  LOG_UNAVAILABLE_HTTP_ERROR = 400
+
+  attr_reader :uri, :bearer_token, :namespace
+
+  def initialize(uri:, bearer_token:, namespace:)
+    @uri = uri
+    @bearer_token = bearer_token
+    @namespace = namespace
+  end
+
+  def api_info
+    handle_exception { pod_client.api }
+  end
+
+  # rubocop:disable Metrics/ParameterLists
+  def create_pod(pod_name:, image:, cmd:, internal_mounts: [], external_mounts: [], node_selector:)
+    handle_exception do
+      pod = Kubeclient::Resource.new(
+        metadata: {
+          name: pod_name,
+          namespace: namespace
+        },
+        spec: {
+          containers: [
+            container_options(name: pod_name, image: image, cmd: cmd, internal_mounts: internal_mounts)
+          ],
+          restartPolicy: "Never",
+          nodeSelector: { node_selector => "" },
+          tolerations: [
+            {
+              key: node_selector,
+              effect: "NoSchedule"
+            }
+          ],
+          volumes: external_mounts
+        }
+      )
+
+      pod_client.create_pod(pod)
+
+      pod_name
+    end
+  end
+  # rubocop:enable Metrics/ParameterLists
+
+  def fetch_pod_logs(pod_name:)
+    handle_exception(pod_name) do
+      pod_client.get_pod_log(pod_name, namespace).body
+    rescue Kubeclient::HttpError => e
+      raise LogsNotFoundError if e.error_code == LOG_UNAVAILABLE_HTTP_ERROR
+
+      raise
+    end
+  end
+
+  def fetch_pod(pod_name:)
+    handle_exception(pod_name) { pod_client.get_pod(pod_name, namespace) }
+  end
+
+  def force_delete_pod(pod_name:)
+    handle_exception(pod_name) { pod_client.delete_pod(pod_name, namespace, delete_options: delete_options) }
+  end
+
+  def fetch_pods
+    handle_exception do
+      pod_client
+        .get_pods(namespace: namespace)
+        .each_with_object({}) do |pod, result|
+          result[pod.metadata.name] = pod
+        end
+    end
+  end
+
+  def handle_exception(pod_name = nil)
+    yield
+  rescue Kubeclient::ResourceNotFoundError
+    raise PodNotFoundError, "Pod not found #{pod_name}"
+  rescue Kubeclient::HttpError, SocketError, Errno::ECONNREFUSED, OpenSSL::SSL::SSLError => e
+    raise NetworkError, "#{e.class}: #{e.message}"
+  end
+
+  private
+
+  def container_options(name:, image:, cmd:, internal_mounts:)
+    {
+      name: name,
+      image: image,
+      command: ["sh", "-c", cmd],
+      resources: {
+        requests: { cpu: Settings.kubernetes.requests.cpu, memory: Settings.kubernetes.requests.memory },
+        limits: { cpu: Settings.kubernetes.limits.cpu, memory: Settings.kubernetes.limits.memory }
+      },
+      volumeMounts: internal_mounts,
+      securityContext: {
+        runAsUser: Settings.run_container_as.user_id,
+        runAsGroup: Settings.run_container_as.group_id
+      },
+      livenessProbe: liveness_probe_options(internal_mounts: internal_mounts)
+    }
+  end
+
+  def liveness_probe_options(internal_mounts:)
+    return if internal_mounts.empty?
+
+    {
+      exec: {
+        command: %w[ls] + internal_mounts.map { |mount| mount[:mountPath] }
+      },
+      periodSeconds: Settings.kubernetes.liveness_probe_seconds_interval
+    }
+  end
+
+  def delete_options
+    Kubeclient::Resource.new(
+      apiVersion: "v1",
+      gracePeriodSeconds: 0,
+      kind: "DeleteOptions"
+    )
+  end
+
+  def pod_client
+    Kubeclient::Client.new(build_client_uri(path: "/api"), "v1", auth_options: { bearer_token: bearer_token }, ssl_options: { verify_ssl: false })
+  end
+
+  def build_client_uri(path:)
+    parsed_uri = URI.parse(uri)
+    URI::Generic.build(host: parsed_uri.host, port: parsed_uri.port, scheme: parsed_uri.scheme, path: path).to_s
+  end
+end

data/app/services/least_used_node.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+class LeastUsedNode
+  attr_reader :execution_type
+
+  def initialize(execution_type:)
+    @execution_type = execution_type
+  end
+
+  def call
+    least_used_nodes&.sample
+  end
+
+  private
+
+  def least_used_nodes
+    nodes_by_usage[nodes_by_usage.keys.min]
+  end
+
+  def nodes_by_usage
+    @nodes_by_usage ||=
+      nodes
+      .filter { |node| slots(node.id).to_a.filter(&:available?).any? }
+      .group_by { |node| SlotsUsagePercentage.new(slots(node.id)).perform }
+  end
+
+  def nodes
+    @nodes ||=
+      Node
+      .accepting_new_tasks
+      .available
+  end
+
+  def slots(node_id)
+    @slots ||=
+      Slot
+      .only(:id, :node_id, :status)
+      .where(:node_id.in => nodes.map(&:id))
+      .where(execution_type: execution_type)
+      .group_by(&:node_id)
+
+    @slots[node_id]
+  end
+end

data/app/services/lock_manager.rb

@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+class LockManager
+  attr_reader :expire, :wait, :locked, :key
+  KEY_PREFIX = "lockmanager"
+
+  def initialize(type:, id:, expire:, wait: true)
+    @key = "#{KEY_PREFIX}-#{type}-#{id}"
+    @expire = expire
+    @wait = wait
+  end
+
+  def lock
+    if lock!
+      begin
+        yield(self)
+      ensure
+        unlock!
+      end
+    else
+      false
+    end
+  end
+
+  def lock!
+    try_lock
+
+    if wait
+      until locked
+        sleep 0.1
+        try_lock
+      end
+    end
+
+    locked
+  end
+
+  def unlock!
+    redis_client.del(key)
+    @locked = false
+  end
+
+  def keep_locked
+    raise "Lock not acquired" unless locked
+
+    if redis_set(xx: true)
+      puts "[LockManager] lock extended by #{expire}"
+    else
+      raise "[LockManager] Lock expired"
+    end
+  end
+
+  def self.active_locks
+    redis_client.keys("#{KEY_PREFIX}*").each_with_object({}) do |key, result|
+      result[key] = redis_client.ttl(key)
+    end
+  end
+
+  def try_lock
+    @locked = redis_set(nx: true)
+  end
+
+  def redis_set(options)
+    redis_client.set(key, 1, **options.merge(ex: expire))
+  end
+
+  def redis_client
+    LockManager.redis_client
+  end
+
+  def self.redis_client
+    Redis.new(RedisUrlParser.call(Settings.redis_url))
+  end
+end

data/app/services/lock_slot.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+class LockSlot
+  attr_reader :execution_type, :node
+
+  def initialize(execution_type:, node: nil)
+    @execution_type = execution_type
+    @node = node
+  end
+
+  def perform
+    return unless selected_node
+
+    selected_node
+      .slots
+      .available
+      .where(execution_type: execution_type)
+      .find_one_and_update(
+        {
+          "$set" => {
+            status: "attaching"
+          }
+        },
+        return_document: :after
+      )
+  end
+
+  private
+
+  def selected_node
+    @selected_node ||= node || least_used_node
+  end
+
+  def least_used_node
+    LeastUsedNode.new(execution_type: execution_type).call
+  end
+end