solid_mcp 0.2.2 → 0.5.0

data/ext/solid_mcp_native/core/src/writer.rs

@@ -0,0 +1,252 @@
+//! Async message writer with batching
+//!
+//! Uses Tokio channels for non-blocking enqueue and background batch writes.
+
+use crate::db::{Database, DbPool};
+use crate::{Config, Error, Message, Result};
+use std::sync::Arc;
+use tokio::sync::mpsc;
+use tokio::task::JoinHandle;
+use tracing::{debug, error, info, warn};
+
+/// Message writer that batches writes to the database
+pub struct MessageWriter {
+    tx: mpsc::Sender<WriterCommand>,
+    handle: JoinHandle<()>,
+}
+
+enum WriterCommand {
+    Message(Message),
+    Flush(tokio::sync::oneshot::Sender<()>),
+    Shutdown,
+}
+
+impl MessageWriter {
+    /// Create a new message writer
+    pub async fn new(db: Arc<DbPool>, config: &Config) -> Result<Self> {
+        let (tx, rx) = mpsc::channel(config.max_queue_size);
+        let batch_size = config.batch_size;
+        let _shutdown_timeout = config.shutdown_timeout; // TODO: Use for timeout handling
+
+        let handle = tokio::spawn(async move {
+            writer_loop(rx, db, batch_size).await;
+            debug!("MessageWriter worker shutdown complete");
+        });
+
+        info!(
+            "MessageWriter started with batch_size={}, queue_size={}",
+            batch_size, config.max_queue_size
+        );
+
+        Ok(Self { tx, handle })
+    }
+
+    /// Enqueue a message for writing (non-blocking)
+    ///
+    /// Returns `Ok(true)` if enqueued, `Ok(false)` if queue is full.
+    pub fn enqueue(&self, message: Message) -> Result<bool> {
+        match self.tx.try_send(WriterCommand::Message(message)) {
+            Ok(()) => Ok(true),
+            Err(mpsc::error::TrySendError::Full(_)) => {
+                warn!("MessageWriter queue full, dropping message");
+                Ok(false)
+            }
+            Err(mpsc::error::TrySendError::Closed(_)) => Err(Error::Shutdown),
+        }
+    }
+
+    /// Enqueue a message for writing (async, waits if queue is full)
+    pub async fn enqueue_async(&self, message: Message) -> Result<()> {
+        self.tx
+            .send(WriterCommand::Message(message))
+            .await
+            .map_err(|_| Error::Shutdown)
+    }
+
+    /// Flush all pending messages to the database
+    pub async fn flush(&self) -> Result<()> {
+        let (tx, rx) = tokio::sync::oneshot::channel();
+        self.tx
+            .send(WriterCommand::Flush(tx))
+            .await
+            .map_err(|_| Error::Shutdown)?;
+        rx.await.map_err(|_| Error::Shutdown)
+    }
+
+    /// Shutdown the writer gracefully
+    pub async fn shutdown(self) -> Result<()> {
+        info!("MessageWriter shutting down...");
+
+        // Send shutdown command
+        let _ = self.tx.send(WriterCommand::Shutdown).await;
+
+        // Wait for worker to finish
+        self.handle
+            .await
+            .map_err(|e| Error::Config(format!("Worker panicked: {}", e)))?;
+
+        info!("MessageWriter shutdown complete");
+        Ok(())
+    }
+}
+
+async fn writer_loop(mut rx: mpsc::Receiver<WriterCommand>, db: Arc<DbPool>, batch_size: usize) {
+    let mut batch = Vec::with_capacity(batch_size);
+    let mut flush_waiters: Vec<tokio::sync::oneshot::Sender<()>> = Vec::new();
+
+    loop {
+        // Wait for first message or command
+        let cmd = match rx.recv().await {
+            Some(cmd) => cmd,
+            None => {
+                debug!("Channel closed, exiting writer loop");
+                break;
+            }
+        };
+
+        match cmd {
+            WriterCommand::Message(msg) => {
+                batch.push(msg);
+            }
+            WriterCommand::Flush(waiter) => {
+                flush_waiters.push(waiter);
+            }
+            WriterCommand::Shutdown => {
+                debug!("Shutdown command received");
+                // Drain remaining messages
+                drain_remaining(&mut rx, &mut batch, &mut flush_waiters);
+                // Write final batch
+                if !batch.is_empty() {
+                    write_batch(&db, &mut batch).await;
+                }
+                // Signal all flush waiters
+                signal_flush_waiters(&mut flush_waiters);
+                break;
+            }
+        }
+
+        // Try to fill batch (non-blocking)
+        while batch.len() < batch_size {
+            match rx.try_recv() {
+                Ok(WriterCommand::Message(msg)) => {
+                    batch.push(msg);
+                }
+                Ok(WriterCommand::Flush(waiter)) => {
+                    flush_waiters.push(waiter);
+                    break; // Stop filling, write now
+                }
+                Ok(WriterCommand::Shutdown) => {
+                    drain_remaining(&mut rx, &mut batch, &mut flush_waiters);
+                    if !batch.is_empty() {
+                        write_batch(&db, &mut batch).await;
+                    }
+                    signal_flush_waiters(&mut flush_waiters);
+                    return;
+                }
+                Err(_) => break, // No more messages available
+            }
+        }
+
+        // Write batch if non-empty
+        if !batch.is_empty() {
+            write_batch(&db, &mut batch).await;
+        }
+
+        // Signal flush waiters
+        signal_flush_waiters(&mut flush_waiters);
+    }
+}
+
+fn drain_remaining(
+    rx: &mut mpsc::Receiver<WriterCommand>,
+    batch: &mut Vec<Message>,
+    flush_waiters: &mut Vec<tokio::sync::oneshot::Sender<()>>,
+) {
+    while let Ok(cmd) = rx.try_recv() {
+        match cmd {
+            WriterCommand::Message(msg) => batch.push(msg),
+            WriterCommand::Flush(waiter) => flush_waiters.push(waiter),
+            WriterCommand::Shutdown => {}
+        }
+    }
+}
+
+async fn write_batch(db: &DbPool, batch: &mut Vec<Message>) {
+    let count = batch.len();
+    debug!("Writing batch of {} messages", count);
+
+    match db.insert_batch(batch).await {
+        Ok(()) => {
+            debug!("Successfully wrote {} messages", count);
+        }
+        Err(e) => {
+            error!("Failed to write batch: {}", e);
+            // TODO: Implement retry logic or dead letter queue
+        }
+    }
+
+    batch.clear();
+}
+
+fn signal_flush_waiters(waiters: &mut Vec<tokio::sync::oneshot::Sender<()>>) {
+    for waiter in waiters.drain(..) {
+        let _ = waiter.send(());
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::db::sqlite::SqlitePool;
+
+    async fn create_test_db() -> Arc<DbPool> {
+        let sqlite = SqlitePool::new("sqlite::memory:").await.unwrap();
+        sqlite.setup_test_schema().await.unwrap();
+        Arc::new(DbPool::Sqlite(sqlite))
+    }
+
+    #[tokio::test]
+    async fn test_writer_basic() {
+        let db = create_test_db().await;
+        let config = Config::new("sqlite::memory:").batch_size(10);
+
+        let writer = MessageWriter::new(db.clone(), &config).await.unwrap();
+
+        // Enqueue some messages
+        for i in 0..5 {
+            let msg = Message::new("session-1", "message", format!(r#"{{"i":{}}}"#, i));
+            assert!(writer.enqueue(msg).unwrap());
+        }
+
+        // Flush to ensure writes complete
+        writer.flush().await.unwrap();
+
+        // Verify messages in database
+        let messages = db.fetch_after("session-1", 0, 100).await.unwrap();
+        assert_eq!(messages.len(), 5);
+
+        // Shutdown
+        writer.shutdown().await.unwrap();
+    }
+
+    #[tokio::test]
+    async fn test_writer_batching() {
+        let db = create_test_db().await;
+        let config = Config::new("sqlite::memory:").batch_size(3);
+
+        let writer = MessageWriter::new(db.clone(), &config).await.unwrap();
+
+        // Enqueue more messages than batch size
+        for i in 0..10 {
+            let msg = Message::new("session-1", "message", format!(r#"{{"i":{}}}"#, i));
+            writer.enqueue_async(msg).await.unwrap();
+        }
+
+        writer.flush().await.unwrap();
+
+        let messages = db.fetch_after("session-1", 0, 100).await.unwrap();
+        assert_eq!(messages.len(), 10);
+
+        writer.shutdown().await.unwrap();
+    }
+}

data/ext/solid_mcp_native/extconf.rb

@@ -0,0 +1,3 @@
+# frozen_string_literal: true
+
+require_relative 'ffi/extconf'

data/ext/solid_mcp_native/ffi/Cargo.toml

@@ -0,0 +1,20 @@
+[package]
+name = "solid_mcp_native"
+version = "0.1.0"
+edition = "2024"
+rust-version = "1.85"
+authors = ["Abdelkader Boudih <terminale@gmail.com>"]
+license = "MIT"
+
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+solid-mcp-core = { workspace = true }
+magnus = { version = "0.8", features = ["embed"] }
+tokio = { workspace = true }
+tracing = { workspace = true }
+tracing-subscriber = { version = "0.3", features = ["env-filter"] }
+
+[build-dependencies]
+rb-sys = "0.9"

data/ext/solid_mcp_native/ffi/extconf.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+def create_noop_makefile(message)
+  warn message
+  warn 'SolidMCP will fall back to pure Ruby backend.'
+  File.write('Makefile', <<~MAKE)
+    all:
+    	@echo '#{message}'
+    install:
+    	@echo '#{message}'
+  MAKE
+  exit 0
+end
+
+# Skip native extension compilation on JRuby
+if RUBY_ENGINE == 'jruby'
+  create_noop_makefile('Skipping native extension on JRuby')
+end
+
+# TruffleRuby 24.0.0+ has native C extension support
+if RUBY_ENGINE == 'truffleruby'
+  warn '⚠️  TruffleRuby detected - C extension support is experimental'
+  warn '    Attempting compilation... (may fail, will fall back to pure Ruby)'
+end
+
+# Check if Cargo is available
+def cargo_available?
+  system('cargo --version > /dev/null 2>&1')
+end
+
+unless cargo_available?
+  create_noop_makefile('Skipping native extension (Cargo not found)')
+end
+
+# Use rb_sys to compile the Rust extension
+require 'mkmf'
+
+# Wrap entire compilation process in error handling to ensure gem install never fails
+begin
+  require 'rb_sys/mkmf'
+  require 'pathname'
+
+  create_rust_makefile('solid_mcp_native/solid_mcp_native') do |r|
+    ffi_dir = Pathname(__dir__)
+    r.ext_dir = begin
+      ffi_dir.relative_path_from(Pathname(Dir.pwd)).to_s
+    rescue ArgumentError
+      ffi_dir.expand_path.to_s
+    end
+    # Profile configuration
+    r.profile = ENV.fetch('RB_SYS_CARGO_PROFILE', :release).to_sym
+  end
+
+  makefile_path = File.join(Dir.pwd, 'Makefile')
+  if File.exist?(makefile_path)
+    manifest_path = File.expand_path(__dir__)
+    contents = File.read(makefile_path)
+    contents.gsub!(/^RB_SYS_CARGO_MANIFEST_DIR \?=.*$/, "RB_SYS_CARGO_MANIFEST_DIR ?= #{manifest_path}")
+    File.write(makefile_path, contents)
+  end
+rescue LoadError => e
+  # rb_sys not available
+  create_noop_makefile("Skipping native extension (rb_sys gem not available: #{e.message})")
+rescue StandardError => e
+  # Any other compilation setup failure (Rust compilation errors, Makefile generation, etc.)
+  create_noop_makefile("Skipping native extension (compilation setup failed: #{e.message})")
+end

data/ext/solid_mcp_native/ffi/src/lib.rs

@@ -0,0 +1,224 @@
+//! Ruby FFI bridge for solid-mcp-core
+//!
+//! Exposes the Rust pub/sub engine to Ruby via Magnus.
+
+use magnus::{Error, Ruby, function};
+use solid_mcp_core::{Config, PubSub};
+use std::cell::RefCell;
+use std::sync::{Arc, OnceLock};
+use std::time::Duration;
+use tokio::runtime::Runtime;
+use tracing::Level;
+use tracing_subscriber::FmtSubscriber;
+
+// Global Tokio runtime
+static RUNTIME: OnceLock<Runtime> = OnceLock::new();
+
+// Thread-local PubSub instance (each Ruby thread gets its own)
+thread_local! {
+    static PUBSUB: RefCell<Option<Arc<PubSub>>> = const { RefCell::new(None) };
+}
+
+fn get_runtime() -> &'static Runtime {
+    RUNTIME.get_or_init(|| {
+        tokio::runtime::Builder::new_multi_thread()
+            .enable_all()
+            .worker_threads(4)
+            .thread_name("solid-mcp-worker")
+            .build()
+            .expect("Failed to create Tokio runtime")
+    })
+}
+
+/// Helper to create a runtime error
+fn runtime_error(msg: impl Into<String>) -> Error {
+    Error::new(Ruby::get().unwrap().exception_runtime_error(), msg.into())
+}
+
+/// Initialize the pub/sub engine with a database URL
+fn init_engine(database_url: String) -> Result<bool, Error> {
+    // Initialize tracing if DEBUG env var is set
+    if std::env::var("DEBUG_SOLID_MCP").is_ok() {
+        let subscriber = FmtSubscriber::builder()
+            .with_max_level(Level::DEBUG)
+            .finish();
+        let _ = tracing::subscriber::set_global_default(subscriber);
+    }
+
+    let rt = get_runtime();
+
+    let config = Config::new(&database_url);
+
+    let pubsub = rt
+        .block_on(async { PubSub::new(config).await })
+        .map_err(|e| runtime_error(e.to_string()))?;
+
+    PUBSUB.with(|ps| {
+        *ps.borrow_mut() = Some(Arc::new(pubsub));
+    });
+
+    Ok(true)
+}
+
+/// Initialize with custom configuration
+fn init_engine_with_config(
+    database_url: String,
+    batch_size: usize,
+    polling_interval_ms: u64,
+    max_queue_size: usize,
+) -> Result<bool, Error> {
+    let rt = get_runtime();
+
+    let config = Config::new(&database_url)
+        .batch_size(batch_size)
+        .polling_interval(Duration::from_millis(polling_interval_ms))
+        .max_queue_size(max_queue_size);
+
+    let pubsub = rt
+        .block_on(async { PubSub::new(config).await })
+        .map_err(|e| runtime_error(e.to_string()))?;
+
+    PUBSUB.with(|ps| {
+        *ps.borrow_mut() = Some(Arc::new(pubsub));
+    });
+
+    Ok(true)
+}
+
+/// Broadcast a message to a session (non-blocking)
+fn broadcast(session_id: String, event_type: String, data: String) -> Result<bool, Error> {
+    PUBSUB.with(|ps| {
+        let ps = ps.borrow();
+        let pubsub = ps.as_ref().ok_or_else(|| {
+            runtime_error("Engine not initialized")
+        })?;
+
+        pubsub
+            .broadcast(&session_id, &event_type, &data)
+            .map_err(|e| runtime_error(e.to_string()))
+    })
+}
+
+/// Flush all pending messages to the database
+fn flush() -> Result<bool, Error> {
+    let rt = get_runtime();
+
+    PUBSUB.with(|ps| {
+        let ps = ps.borrow();
+        let pubsub = ps.as_ref().ok_or_else(|| {
+            runtime_error("Engine not initialized")
+        })?;
+
+        rt.block_on(async { pubsub.flush().await })
+            .map_err(|e| runtime_error(e.to_string()))?;
+
+        Ok(true)
+    })
+}
+
+/// Mark messages as delivered
+fn mark_delivered(ids: Vec<i64>) -> Result<bool, Error> {
+    let rt = get_runtime();
+
+    PUBSUB.with(|ps| {
+        let ps = ps.borrow();
+        let pubsub = ps.as_ref().ok_or_else(|| {
+            runtime_error("Engine not initialized")
+        })?;
+
+        rt.block_on(async { pubsub.mark_delivered(&ids).await })
+            .map_err(|e| runtime_error(e.to_string()))?;
+
+        Ok(true)
+    })
+}
+
+/// Cleanup old messages
+/// Returns [delivered_count, undelivered_count]
+fn cleanup() -> Result<Vec<u64>, Error> {
+    let rt = get_runtime();
+
+    PUBSUB.with(|ps| {
+        let ps = ps.borrow();
+        let pubsub = ps.as_ref().ok_or_else(|| {
+            runtime_error("Engine not initialized")
+        })?;
+
+        let (delivered, undelivered) = rt
+            .block_on(async { pubsub.cleanup().await })
+            .map_err(|e| runtime_error(e.to_string()))?;
+
+        Ok(vec![delivered, undelivered])
+    })
+}
+
+/// Shutdown the pub/sub engine
+fn shutdown() -> Result<bool, Error> {
+    let rt = get_runtime();
+
+    PUBSUB.with(|ps| {
+        let mut ps = ps.borrow_mut();
+        if let Some(pubsub) = ps.take() {
+            // Try to get exclusive access
+            match Arc::try_unwrap(pubsub) {
+                Ok(pubsub) => {
+                    let _ = rt.block_on(async { pubsub.shutdown().await });
+                }
+                Err(_) => {
+                    // Other references exist, can't fully shutdown
+                    tracing::warn!("Cannot fully shutdown: other references exist");
+                }
+            }
+        }
+        Ok(true)
+    })
+}
+
+/// Get the library version
+fn version() -> &'static str {
+    env!("CARGO_PKG_VERSION")
+}
+
+/// Check if the engine is initialized
+fn initialized() -> bool {
+    PUBSUB.with(|ps| ps.borrow().is_some())
+}
+
+/// Get subscription count
+fn subscription_count() -> Result<usize, Error> {
+    let rt = get_runtime();
+
+    PUBSUB.with(|ps| {
+        let ps = ps.borrow();
+        let pubsub = ps.as_ref().ok_or_else(|| {
+            runtime_error("Engine not initialized")
+        })?;
+
+        Ok(rt.block_on(async { pubsub.subscription_count().await }))
+    })
+}
+
+#[magnus::init]
+fn init(ruby: &Ruby) -> Result<(), Error> {
+    let module = ruby.define_module("SolidMCPNative")?;
+
+    // Core functions
+    module.define_module_function("version", function!(version, 0))?;
+    module.define_module_function("initialized?", function!(initialized, 0))?;
+
+    // Lifecycle
+    module.define_module_function("init", function!(init_engine, 1))?;
+    module.define_module_function("init_with_config", function!(init_engine_with_config, 4))?;
+    module.define_module_function("shutdown", function!(shutdown, 0))?;
+
+    // Messaging
+    module.define_module_function("broadcast", function!(broadcast, 3))?;
+    module.define_module_function("flush", function!(flush, 0))?;
+    module.define_module_function("mark_delivered", function!(mark_delivered, 1))?;
+    module.define_module_function("cleanup", function!(cleanup, 0))?;
+
+    // Status
+    module.define_module_function("subscription_count", function!(subscription_count, 0))?;
+
+    Ok(())
+}

data/lib/solid_mcp/configuration.rb

@@ -2,8 +2,9 @@
 
 module SolidMCP
   class Configuration
-    attr_accessor :batch_size, :flush_interval, :delivered_retention, 
-                  :undelivered_retention, :polling_interval, :max_wait_time, :logger
+    attr_accessor :batch_size, :flush_interval, :delivered_retention,
+                  :undelivered_retention, :polling_interval, :max_wait_time, :logger,
+                  :max_queue_size, :shutdown_timeout
 
     def initialize
       @batch_size = 200
@@ -12,6 +13,8 @@ module SolidMCP
       @max_wait_time = 30 # 30 seconds
       @delivered_retention = 3600 # 1 hour in seconds
       @undelivered_retention = 86400 # 24 hours in seconds
+      @max_queue_size = 10_000 # Maximum messages in memory queue
+      @shutdown_timeout = 30 # Maximum seconds to wait for graceful shutdown
       @logger = default_logger
     end