itsi 0.1.0 → 0.1.3

data/gems/scheduler/ext/itsi_server/src/server/tls.rs

@@ -0,0 +1,152 @@
+use base64::{engine::general_purpose, Engine as _};
+use itsi_error::Result;
+use itsi_tracing::{info, warn};
+use rcgen::{CertificateParams, DnType, KeyPair, SanType};
+use rustls_pemfile::{certs, pkcs8_private_keys};
+use std::{collections::HashMap, fs, io::BufReader};
+use tokio_rustls::rustls::{Certificate, PrivateKey, ServerConfig};
+
+const ITS_CA_CERT: &str = include_str!("./itsi_ca/itsi_ca.crt");
+const ITS_CA_KEY: &str = include_str!("./itsi_ca/itsi_ca.key");
+
+// Generates a TLS configuration based on either :
+// * Input "cert" and "key" options (either paths or Base64-encoded strings) or
+// * Performs automatic certificate generation/retrieval. Generated certs use an internal self-signed Isti CA.
+// If a non-local host or optional domain parameter is provided,
+// an automated certificate will attempt to be fetched using let's encrypt.
+pub fn configure_tls(host: &str, query_params: &HashMap<String, String>) -> Result<ServerConfig> {
+    info!("TLS Options {:?}", query_params);
+    let (certs, key) = if let (Some(cert_path), Some(key_path)) =
+        (query_params.get("cert"), query_params.get("key"))
+    {
+        // Load from file or Base64
+        let certs = load_certs(cert_path);
+        let key = load_private_key(key_path);
+        (certs, key)
+    } else {
+        let domains_param = query_params
+            .get("domains")
+            .map(|v| v.split(',').map(String::from).collect());
+        let host_string = host.to_string();
+        let domains = domains_param.or_else(|| {
+            if host_string != "localhost" {
+                Some(vec![host_string])
+            } else {
+                None
+            }
+        });
+
+        if let Some(domains) = domains {
+            retrieve_acme_cert(domains)?
+        } else {
+            generate_ca_signed_cert(vec![host.to_owned()])?
+        }
+    };
+
+    let mut config = ServerConfig::builder()
+        .with_safe_defaults()
+        .with_no_client_auth()
+        .with_single_cert(certs, key)
+        .expect("Failed to build TLS config");
+
+    config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
+    Ok(config)
+}
+
+pub fn load_certs(path: &str) -> Vec<Certificate> {
+    let data = if let Some(stripped) = path.strip_prefix("base64:") {
+        general_purpose::STANDARD
+            .decode(stripped)
+            .expect("Invalid base64 certificate")
+    } else {
+        fs::read(path).expect("Failed to read certificate file")
+    };
+
+    if data.starts_with(b"-----BEGIN ") {
+        let mut reader = BufReader::new(&data[..]);
+        let certs_der: Vec<Vec<u8>> = certs(&mut reader)
+            .map(|r| {
+                r.map(|der| der.as_ref().to_vec())
+                    .map_err(itsi_error::ItsiError::from)
+            })
+            .collect::<Result<_>>()
+            .expect("Failed to parse certificate file");
+        certs_der.into_iter().map(Certificate).collect()
+    } else {
+        vec![Certificate(data)]
+    }
+}
+
+/// Loads a private key from a file or Base64.
+pub fn load_private_key(path: &str) -> PrivateKey {
+    let key_data = if let Some(stripped) = path.strip_prefix("base64:") {
+        general_purpose::STANDARD
+            .decode(stripped)
+            .expect("Invalid base64 private key")
+    } else {
+        fs::read(path).expect("Failed to read private key file")
+    };
+
+    if key_data.starts_with(b"-----BEGIN ") {
+        let mut reader = BufReader::new(&key_data[..]);
+        let keys: Vec<Vec<u8>> = pkcs8_private_keys(&mut reader)
+            .map(|r| {
+                r.map(|key| key.secret_pkcs8_der().to_vec())
+                    .map_err(itsi_error::ItsiError::from)
+            })
+            .collect::<Result<_>>()
+            .expect("Failed to parse private key");
+        if !keys.is_empty() {
+            return PrivateKey(keys[0].clone());
+        }
+    }
+    PrivateKey(key_data)
+}
+
+pub fn generate_ca_signed_cert(domains: Vec<String>) -> Result<(Vec<Certificate>, PrivateKey)> {
+    info!("Generating New Itsi CA - Self signed Certificate. Use `itsi ca export` to export the CA certificate for import into your local trust store.");
+
+    let ca_kp = KeyPair::from_pem(ITS_CA_KEY).expect("Failed to load embedded CA key");
+    let ca_cert = CertificateParams::from_ca_cert_pem(ITS_CA_CERT)
+        .expect("Failed to parse embedded CA certificate")
+        .self_signed(&ca_kp)
+        .expect("Failed to self-sign embedded CA cert");
+
+    let ee_key = KeyPair::generate_for(&rcgen::PKCS_ECDSA_P256_SHA256).unwrap();
+    let mut ee_params = CertificateParams::default();
+
+    info!(
+        "Generated certificate will be valid for domains {:?}",
+        domains
+    );
+    use std::net::IpAddr;
+
+    ee_params.subject_alt_names = domains
+        .iter()
+        .map(|domain| {
+            if let Ok(ip) = domain.parse::<IpAddr>() {
+                SanType::IpAddress(ip)
+            } else {
+                SanType::DnsName(domain.clone().try_into().unwrap())
+            }
+        })
+        .collect();
+
+    ee_params
+        .distinguished_name
+        .push(DnType::CommonName, domains[0].clone());
+
+    ee_params.use_authority_key_identifier_extension = true;
+
+    let ee_cert = ee_params.signed_by(&ee_key, &ca_cert, &ca_kp).unwrap();
+    let ee_cert_der = ee_cert.der().to_vec();
+    let ee_cert = Certificate(ee_cert_der);
+    let ca_cert = Certificate(ca_cert.der().to_vec());
+    Ok((vec![ee_cert, ca_cert], PrivateKey(ee_key.serialize_der())))
+}
+
+/// TODO: Retrieves an ACME certificate for a given domain.
+pub fn retrieve_acme_cert(domains: Vec<String>) -> Result<(Vec<Certificate>, PrivateKey)> {
+    warn!("Retrieving ACME cert for {}", domains.join(", "));
+    generate_ca_signed_cert(domains)
+}

data/gems/scheduler/ext/itsi_tracing/Cargo.toml

@@ -9,4 +9,8 @@ tracing-subscriber = { version = "0.3.19", features = [
   "env-filter",
   "std",
   "fmt",
+  "json",
+  "ansi",
 ] }
+tracing-attributes = "0.1"
+atty = "0.2.14"

data/gems/scheduler/ext/itsi_tracing/src/lib.rs

@@ -1,11 +1,41 @@
+use std::env;
+
+use atty::{Stream, is};
 pub use tracing::{debug, error, info, trace, warn};
-use tracing_subscriber::{EnvFilter, fmt};
+pub use tracing_attributes::instrument; // Explicitly export from tracing-attributes
+use tracing_subscriber::{
+    EnvFilter,
+    fmt::{self, format},
+};
 
+#[instrument]
 pub fn init() {
-    let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
-    let format = fmt::format().with_level(true).with_target(false).compact();
-    tracing_subscriber::fmt()
-        .with_env_filter(env_filter)
+    let env_filter = EnvFilter::builder()
+        .with_env_var("ITSI_LOG")
+        .try_from_env()
+        .unwrap_or_else(|_| EnvFilter::new("info"));
+
+    let format = fmt::format()
+        .compact()
+        .with_file(false)
+        .with_level(true)
+        .with_line_number(false)
+        .with_source_location(false)
+        .with_target(false)
+        .with_thread_ids(false);
+
+    let is_tty = is(Stream::Stdout);
+
+    let subscriber = tracing_subscriber::fmt()
         .event_format(format)
-        .init();
+        .with_env_filter(env_filter);
+
+    if (is_tty && env::var("ITSI_LOG_PLAIN").is_err()) || env::var("ITSI_LOG_ANSI").is_ok() {
+        subscriber.with_ansi(true).init();
+    } else {
+        subscriber
+            .fmt_fields(format::JsonFields::default())
+            .event_format(fmt::format().json())
+            .init();
+    }
 }

data/gems/scheduler/itsi-scheduler.gemspec

@@ -12,8 +12,8 @@ Gem::Specification.new do |spec|
   spec.description = "Itsi Scheduler - A light-weight Fiber Scheduler implementation for Ruby"
   spec.homepage = "https://itsi.fyi"
   spec.license = "MIT"
-  spec.required_ruby_version = ">= 3.1.0"
-  spec.required_rubygems_version = ">= 3.3.11"
+  spec.required_ruby_version = ">= 3.0.0"
+  spec.required_rubygems_version = ">= 3.1.11"
 
   spec.metadata["homepage_uri"] = spec.homepage
   spec.metadata["source_code_uri"] = "https://github.com/wouterken/itsi/scheduler"
@@ -35,7 +35,6 @@ Gem::Specification.new do |spec|
 
   # Uncomment to register a new dependency of your gem
   # spec.add_dependency "example-gem", "~> 1.0"
-  spec.add_dependency "libclang", "~> 14.0"
   spec.add_dependency "rb_sys", "~> 0.9.91"
 
   # For more information and examples about making a new gem, check out our

data/gems/scheduler/lib/itsi/scheduler/version.rb

@@ -2,6 +2,6 @@
 
 module Itsi
   class Scheduler
-    VERSION = "0.1.0"
+    VERSION = "0.1.3"
   end
 end

data/gems/scheduler/lib/itsi/scheduler.rb

@@ -6,6 +6,142 @@ require_relative "scheduler/itsi_scheduler"
 module Itsi
   class Scheduler
     class Error < StandardError; end
-    # Your code goes here...
+
+    def self.resume_token
+      @resume_token ||= 0
+      @resume_token += 1
+    end
+
+    def initialize
+      @join_waiters = {}.compare_by_identity
+      @token_map = {}.compare_by_identity
+      @resume_tokens = {}.compare_by_identity
+      @unblocked = [[], []]
+      @unblock_idx = 0
+      @unblocked_mux = Mutex.new
+      @resume_fiber = method(:resume_fiber).to_proc
+      @resume_fiber_with_readiness = method(:resume_fiber_with_readiness).to_proc
+      @resume_blocked = method(:resume_blocked).to_proc
+    end
+
+    def block(_, timeout, fiber = Fiber.current, token = Scheduler.resume_token)
+      @join_waiters[fiber] = true
+
+      start_timer(timeout, token) if timeout
+      @resume_tokens[token] = fiber
+      @token_map[fiber] = token
+      Fiber.yield
+    ensure
+      @resume_tokens.delete(token)
+      @token_map.delete(fiber)
+      @join_waiters.delete(fiber)
+    end
+
+    # Register an IO waiter.
+    # This will get resumed by our scheduler inside the call to
+    # fetch_events.
+    def io_wait(io, events, duration)
+      fiber = Fiber.current
+      token = Scheduler.resume_token
+      readiness = register_io_wait(io.fileno, events, duration, token)
+      readiness || block(nil, duration, fiber, token)
+    end
+
+    def unblock(_blocker, fiber)
+      @unblocked_mux.synchronize do
+        @unblocked[@unblock_idx] << fiber
+      end
+      wake
+    end
+
+    def kernel_sleep(duration)
+      block nil, duration
+    end
+
+    def tick
+      events = fetch_due_events
+      timers = fetch_due_timers
+      unblocked = switch_unblock_batch
+      events&.each(&@resume_fiber_with_readiness)
+      unblocked.each(&@resume_blocked)
+      unblocked.clear
+      timers&.each(&@resume_fiber)
+    end
+
+    def resume_fiber(token)
+      if (fiber = @resume_tokens.delete(token))
+        fiber.resume
+      end
+    rescue StandardError => e
+      warn "Failed to resume fiber #{fiber}: #{e.message}"
+    end
+
+    def resume_fiber_with_readiness((token, readiness))
+      if (fiber = @resume_tokens.delete(token))
+        fiber.resume(readiness)
+      end
+    rescue StandardError => e
+      warn "Failed to resume fiber #{fiber}: #{e.message}"
+    end
+
+    def resume_blocked(fiber)
+      if (token = @token_map[fiber])
+        resume_fiber(token)
+      elsif fiber.alive?
+        fiber.resume
+      end
+    end
+
+    def switch_unblock_batch
+      @unblocked_mux.synchronize do
+        current = @unblocked[@unblock_idx]
+        @unblock_idx = (@unblock_idx + 1) % 2
+        current
+      end
+    end
+
+    # Yields upwards to the scheduler, with an intention to
+    # resume the fiber that yielded ASAP.
+    def yield
+      kernel_sleep(0) if work?
+    end
+
+    # Keep running until we've got no timers we're awaiting, no pending IO, no temporary yields,
+    # no pending unblocks.
+    def work?
+      !@unblocked[@unblock_idx].empty? || !@join_waiters.empty? || has_pending_io?
+    end
+
+    # Run until no more work needs doing.
+    def run
+      tick while work?
+      debug "Exit Scheduler"
+    end
+
+    # Hook invoked at the end of the thread.
+    # Will start our scheduler's Reactor.
+    def scheduler_close
+      run
+    ensure
+      @closed ||= true
+      freeze
+    end
+
+    # Need to defer to Process::Status rather than our extension
+    # as we don't have a means of creating our own Process::Status.
+    def process_wait(pid, flags)
+      Thread.new do
+        Process::Status.wait(pid, flags)
+      end.value
+    end
+
+    def closed?
+      @closed
+    end
+
+    # Spin up a new fiber and immediately resume it.
+    def fiber(&blk)
+      Fiber.new(blocking: false, &blk).tap(&:resume)
+    end
   end
 end

data/gems/scheduler/test/helpers/test_helper.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+require "minitest/reporters"
+Minitest::Reporters.use! Minitest::Reporters::SpecReporter.new
+
+require "itsi/scheduler"
+require 'debug'
+module Itsi::Scheduler::TestHelper
+  SchedulerClass = Itsi::Scheduler
+
+  def with_scheduler(join: true, report_on_exception: false)
+    Thread.new do
+      Thread.current.report_on_exception = report_on_exception
+      scheduler = SchedulerClass.new
+      Fiber.set_scheduler(scheduler)
+      Fiber.schedule do
+        yield scheduler
+      end
+    end.yield_self do |thread|
+      thread.join if join
+      thread
+    end
+  end
+end

data/gems/scheduler/test/test_active_record.rb

@@ -0,0 +1,158 @@
+# frozen_string_literal: true
+
+require "active_record"
+
+class TestActiveRecordFiberScheduler < Minitest::Test
+  include Itsi::Scheduler::TestHelper
+
+  # Set up an ActiveRecord connection to your PostgreSQL test database.
+  # Adjust the connection parameters as needed.
+  def setup
+    ActiveSupport::IsolatedExecutionState.isolation_level = :fiber
+    ActiveRecord::Base.establish_connection(
+      adapter: "postgresql",
+      database: "fiber_scheduler_test",
+      pool: 2, # use a small pool to test contention scenarios
+      checkout_timeout: 5
+    )
+  end
+
+  # Disconnect after each test.
+  def teardown
+    ActiveRecord::Base.connection_pool.disconnect!
+  end
+
+  # Test a basic query execution inside a fiber.
+  def test_basic_query
+    result = nil
+
+    with_scheduler do |_scheduler|
+      Fiber.schedule do
+        result = ActiveRecord::Base.connection.select_value("SELECT 1")
+      end
+    end
+
+    # select_value returns a string from PG adapter so we compare with "1"
+    assert_equal "1", result.to_s
+  end
+
+  # Test running two queries concurrently in different fibers.
+  def test_concurrent_queries
+    results = []
+
+    with_scheduler do |_scheduler|
+      Fiber.schedule do
+        results << ActiveRecord::Base.connection.select_value("SELECT 1")
+      end
+
+      Fiber.schedule do
+        results << ActiveRecord::Base.connection.select_value("SELECT 2")
+      end
+    end
+
+    # Ensure that both queries have executed and returned the expected values.
+    results = results.map(&:to_s)
+    assert_includes results, "1"
+    assert_includes results, "2"
+  end
+
+  # Test a query that involves a short delay using PostgreSQL's pg_sleep.
+  def test_query_with_delay
+    result = nil
+
+    with_scheduler do |_scheduler|
+      Fiber.schedule do
+        # Introduce a 0.1 second delay.
+        ActiveRecord::Base.connection.execute("SELECT pg_sleep(0.1)")
+        result = ActiveRecord::Base.connection.select_value("SELECT 3")
+      end
+    end
+
+    assert_equal "3", result.to_s
+  end
+
+  # Test connection pool exhaustion by limiting the pool to one connection.
+  # Two fibers will attempt to get a connection concurrently.
+  def test_connection_pool_exhaustion
+    # Re-establish connection with a pool size of 1.
+    ActiveSupport::IsolatedExecutionState.isolation_level = :fiber
+    ActiveRecord::Base.establish_connection(
+      adapter: "postgresql",
+      host: "localhost",
+      database: "fiber_scheduler_test",
+      pool: 1,
+      checkout_timeout: 0.25
+    )
+    # ActiveRecord::Base.connection_pool.disconnect!
+
+    results = []
+
+    with_scheduler do |_scheduler|
+      Fiber.schedule do
+        ActiveRecord::Base.connection_pool.with_connection(prevent_permanent_checkout: true) do
+          results << ActiveRecord::Base.connection.select_value("SELECT 1")
+        end
+      end
+
+      Fiber.schedule do
+        ActiveRecord::Base.connection_pool.with_connection(prevent_permanent_checkout: true) do
+          results << ActiveRecord::Base.connection.select_value("SELECT 2")
+        end
+      end
+    end
+
+    # Takes #{checkout_timeout} seconds between last
+    results = results.map(&:to_s)
+    assert_includes results, "1"
+    assert_includes results, "2"
+  end
+
+  # Test that after a Fiber finishes its work, its connection is automatically released.
+  def test_fiber_connection_release_after_completion
+    # Use the scheduler to run a fiber that checks out a connection and does a simple query.
+
+    with_scheduler do |_scheduler|
+      Fiber.schedule do
+        # This fiber checks out a connection to run a query.
+        ActiveRecord::Base.connection.select_value("SELECT 1")
+        # The fiber ends here.
+      end
+    end
+
+    # After the scheduler finishes, the fiber should have completed and released its connection.
+    # Now we attempt to checkout a connection manually. If the previous fiber's connection
+    # was not released, this would either time out or raise an error.
+    connection = ActiveRecord::Base.connection_pool.checkout
+    assert connection, "Expected to obtain a connection after fiber completion"
+    ActiveRecord::Base.connection_pool.checkin(connection)
+  end
+
+  # Test that a transaction works correctly when run inside a fiber.
+  # A temporary table is created, a record inserted and then queried before the transaction is rolled back.
+  def test_transaction_fiber
+    result = nil
+
+    with_scheduler do |_scheduler|
+      Fiber.schedule do
+        ActiveRecord::Base.transaction do
+          # Create a temporary table for testing.
+          ActiveRecord::Base.connection.execute(<<~SQL)
+            CREATE TEMP TABLE IF NOT EXISTS test_table (
+              id serial PRIMARY KEY,
+              name text
+            )
+          SQL
+
+          # Insert a record.
+          ActiveRecord::Base.connection.execute("INSERT INTO test_table (name) VALUES ('Alice')")
+          # Query the inserted record.
+          result = ActiveRecord::Base.connection.select_value("SELECT name FROM test_table LIMIT 1")
+          # Roll back the transaction to avoid leaving test data.
+          raise ActiveRecord::Rollback
+        end
+      end
+    end
+
+    assert_equal "Alice", result.to_s
+  end
+end

data/gems/scheduler/test/test_address_resolve.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+require 'debug'
+
+
+class TestAddressResolve < Minitest::Test
+  include Itsi::Scheduler::TestHelper
+
+  def test_addess_resolve
+    results = []
+
+    with_scheduler do |_scheduler|
+      Fiber.schedule do
+        results << Addrinfo.getaddrinfo("www.ruby-lang.org", 80, nil, :STREAM)
+      end
+      Fiber.schedule do
+        results << Addrinfo.getaddrinfo("www.google.com", 80, nil, :STREAM)
+      end
+    end
+
+    assert  results.all?{|results| results.find(&:ipv4?) }
+    assert results.all?{|results| results.find(&:ipv6?) }
+  end
+end