Skip to main content

arete_server/
lib.rs

1//! # arete-server
2//!
3//! WebSocket server and projection handlers for Arete streaming pipelines.
4//!
5//! This crate provides a builder API for creating Arete servers that:
6//!
7//! - Process Solana blockchain data via Yellowstone gRPC
8//! - Transform data using the Arete VM
9//! - Stream entity updates over WebSockets to connected clients
10//! - Support multiple streaming modes (State, List, Append)
11//!
12//! ## Quick Start
13//!
14//! ```rust,ignore
15//! use arete_server::{Server, Spec};
16//!
17//! #[tokio::main]
18//! async fn main() -> anyhow::Result<()> {
19//!     Server::builder()
20//!         .spec(my_spec())
21//!         .websocket()
22//!         .bind("[::]:8877".parse()?)
23//!         .health_monitoring()
24//!         .start()
25//!         .await
26//! }
27//! ```
28//!
29//! ## Feature Flags
30//!
31//! - `otel` - OpenTelemetry integration for metrics and distributed tracing
32
33pub mod bus;
34pub mod cache;
35pub mod compression;
36pub mod config;
37pub mod health;
38pub mod http_health;
39pub mod http_server;
40pub mod materialized_view;
41#[cfg(feature = "otel")]
42pub mod metrics;
43pub mod mutation_batch;
44pub mod projector;
45pub mod runtime;
46pub mod sorted_cache;
47pub mod telemetry;
48pub mod view;
49pub mod websocket;
50
51pub use arete_auth::{AsyncVerifier, KeyLoader, Limits, TokenVerifier, VerifyingKey};
52pub use bus::{BusManager, BusMessage};
53pub use cache::{EntityCache, EntityCacheConfig};
54pub use config::{
55    HealthConfig, HttpHealthConfig, HttpServerConfig, ReconnectionConfig, ServerConfig,
56    WebSocketConfig, YellowstoneConfig,
57};
58pub use health::{HealthMonitor, SlotTracker, StreamStatus};
59pub use http_health::HttpHealthServer;
60pub use http_server::HttpServer;
61pub use materialized_view::{MaterializedView, MaterializedViewRegistry, ViewEffect};
62#[cfg(feature = "otel")]
63pub use metrics::Metrics;
64pub use mutation_batch::{EventContext, MutationBatch, SlotContext};
65pub use projector::Projector;
66pub use runtime::Runtime;
67pub use telemetry::{init as init_telemetry, TelemetryConfig};
68#[cfg(feature = "otel")]
69pub use telemetry::{init_with_otel, TelemetryGuard};
70pub use view::{Delivery, Filters, Projection, ViewIndex, ViewSpec};
71pub use websocket::{
72    AllowAllAuthPlugin, AuthContext, AuthDecision, AuthDeny, AuthErrorDetails, ChannelUsageEmitter,
73    ClientInfo, ClientManager, ConnectionAuthRequest, ErrorResponse, Frame, HttpUsageEmitter, Mode,
74    RateLimitConfig, RateLimitResult, RateLimiterConfig, RefreshAuthRequest, RefreshAuthResponse,
75    RetryPolicy, SignedSessionAuthPlugin, SocketIssueMessage, StaticTokenAuthPlugin, Subscription,
76    WebSocketAuthPlugin, WebSocketRateLimiter, WebSocketServer, WebSocketUsageBatch,
77    WebSocketUsageEmitter, WebSocketUsageEnvelope, WebSocketUsageEvent,
78};
79
80use anyhow::Result;
81use arete_interpreter::ast::ViewDef;
82use std::net::SocketAddr;
83use std::sync::Arc;
84
85/// Type alias for a parser setup function.
86pub type ParserSetupFn = Arc<
87    dyn Fn(
88            tokio::sync::mpsc::Sender<MutationBatch>,
89            Option<HealthMonitor>,
90            ReconnectionConfig,
91        ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send>>
92        + Send
93        + Sync,
94>;
95
96pub type ProgramAccountReaderFn =
97    Arc<dyn Fn(&str, &str, &[u8]) -> Result<serde_json::Value> + Send + Sync>;
98
99/// Specification for a Arete server
100/// Contains bytecode, parsers, and program information
101pub struct Spec {
102    pub bytecode: arete_interpreter::compiler::MultiEntityBytecode,
103    pub program_ids: Vec<String>,
104    pub parser_setup: Option<ParserSetupFn>,
105    pub program_account_reader: Option<ProgramAccountReaderFn>,
106    pub entity_specs: Vec<arete_interpreter::ast::SerializableStreamSpec>,
107    pub views: Vec<ViewDef>,
108}
109
110impl Spec {
111    pub fn new(
112        bytecode: arete_interpreter::compiler::MultiEntityBytecode,
113        program_id: impl Into<String>,
114    ) -> Self {
115        Self {
116            bytecode,
117            program_ids: vec![program_id.into()],
118            parser_setup: None,
119            program_account_reader: None,
120            entity_specs: Vec::new(),
121            views: Vec::new(),
122        }
123    }
124
125    pub fn with_parser_setup(mut self, setup_fn: ParserSetupFn) -> Self {
126        self.parser_setup = Some(setup_fn);
127        self
128    }
129
130    pub fn with_program_account_reader(mut self, reader: ProgramAccountReaderFn) -> Self {
131        self.program_account_reader = Some(reader);
132        self
133    }
134
135    pub fn with_entity_specs(
136        mut self,
137        entity_specs: Vec<arete_interpreter::ast::SerializableStreamSpec>,
138    ) -> Self {
139        self.entity_specs = entity_specs;
140        self
141    }
142
143    pub fn with_views(mut self, views: Vec<ViewDef>) -> Self {
144        self.views = views;
145        self
146    }
147}
148
149/// Main server interface with fluent builder API
150pub struct Server;
151
152impl Server {
153    /// Create a new server builder
154    pub fn builder() -> ServerBuilder {
155        ServerBuilder::new()
156    }
157}
158
159/// Builder for configuring and creating a Arete server
160pub struct ServerBuilder {
161    spec: Option<Spec>,
162    views: Option<ViewIndex>,
163    materialized_views: Option<MaterializedViewRegistry>,
164    config: ServerConfig,
165    websocket_auth_plugin: Option<Arc<dyn WebSocketAuthPlugin>>,
166    http_auth_plugin: Option<Arc<dyn WebSocketAuthPlugin>>,
167    websocket_usage_emitter: Option<Arc<dyn WebSocketUsageEmitter>>,
168    websocket_max_clients: Option<usize>,
169    websocket_rate_limit_config: Option<crate::websocket::client_manager::RateLimitConfig>,
170    #[cfg(feature = "otel")]
171    metrics: Option<Arc<Metrics>>,
172}
173
174impl ServerBuilder {
175    fn new() -> Self {
176        Self {
177            spec: None,
178            views: None,
179            materialized_views: None,
180            config: ServerConfig::new(),
181            websocket_auth_plugin: None,
182            http_auth_plugin: None,
183            websocket_usage_emitter: None,
184            websocket_max_clients: None,
185            websocket_rate_limit_config: None,
186            #[cfg(feature = "otel")]
187            metrics: None,
188        }
189    }
190
191    /// Set the specification (bytecode, parsers, program_ids)
192    pub fn spec(mut self, spec: Spec) -> Self {
193        self.spec = Some(spec);
194        self
195    }
196
197    /// Set custom view index
198    pub fn views(mut self, views: ViewIndex) -> Self {
199        self.views = Some(views);
200        self
201    }
202
203    /// Enable metrics collection (requires 'otel' feature)
204    #[cfg(feature = "otel")]
205    pub fn metrics(mut self, metrics: Metrics) -> Self {
206        self.metrics = Some(Arc::new(metrics));
207        self
208    }
209
210    /// Enable WebSocket server with default configuration
211    pub fn websocket(mut self) -> Self {
212        self.config.websocket = Some(WebSocketConfig::default());
213        self
214    }
215
216    /// Configure WebSocket server
217    pub fn websocket_config(mut self, config: WebSocketConfig) -> Self {
218        self.config.websocket = Some(config);
219        self
220    }
221
222    /// Set a WebSocket auth plugin used to authorize inbound connections.
223    pub fn websocket_auth_plugin(mut self, plugin: Arc<dyn WebSocketAuthPlugin>) -> Self {
224        self.websocket_auth_plugin = Some(plugin);
225        self
226    }
227
228    /// Set an HTTP auth plugin used to authorize inbound read requests.
229    pub fn http_auth_plugin(mut self, plugin: Arc<dyn WebSocketAuthPlugin>) -> Self {
230        self.http_auth_plugin = Some(plugin);
231        self
232    }
233
234    /// Set an async usage emitter for billing-grade websocket usage events.
235    pub fn websocket_usage_emitter(mut self, emitter: Arc<dyn WebSocketUsageEmitter>) -> Self {
236        self.websocket_usage_emitter = Some(emitter);
237        self
238    }
239
240    /// Set the maximum number of concurrent WebSocket clients.
241    pub fn websocket_max_clients(mut self, max_clients: usize) -> Self {
242        self.websocket_max_clients = Some(max_clients);
243        self
244    }
245
246    /// Configure rate limiting for WebSocket connections.
247    ///
248    /// This sets global rate limits such as maximum connections per IP,
249    /// timeouts, and rate windows. Per-subject limits are controlled
250    /// via AuthContext.Limits from the authentication token.
251    pub fn websocket_rate_limit_config(
252        mut self,
253        config: crate::websocket::client_manager::RateLimitConfig,
254    ) -> Self {
255        self.websocket_rate_limit_config = Some(config);
256        self
257    }
258
259    /// Set the bind address for WebSocket server
260    pub fn bind(mut self, addr: impl Into<SocketAddr>) -> Self {
261        if let Some(ws_config) = &mut self.config.websocket {
262            ws_config.bind_address = addr.into();
263        } else {
264            self.config.websocket = Some(WebSocketConfig::new(addr.into()));
265        }
266        self
267    }
268
269    /// Configure Yellowstone gRPC connection
270    pub fn yellowstone(mut self, config: YellowstoneConfig) -> Self {
271        self.config.yellowstone = Some(config);
272        self
273    }
274
275    /// Enable health monitoring with default configuration
276    pub fn health_monitoring(mut self) -> Self {
277        self.config.health = Some(HealthConfig::default());
278        self
279    }
280
281    /// Configure health monitoring
282    pub fn health_config(mut self, config: HealthConfig) -> Self {
283        self.config.health = Some(config);
284        self
285    }
286
287    /// Enable reconnection with default configuration
288    pub fn reconnection(mut self) -> Self {
289        self.config.reconnection = Some(ReconnectionConfig::default());
290        self
291    }
292
293    /// Configure reconnection behavior
294    pub fn reconnection_config(mut self, config: ReconnectionConfig) -> Self {
295        self.config.reconnection = Some(config);
296        self
297    }
298
299    /// Enable the HTTP server with default configuration (port 8081).
300    ///
301    /// This serves health endpoints plus stack-scoped HTTP reads.
302    pub fn http(mut self) -> Self {
303        self.config.http_health = Some(HttpHealthConfig::default());
304        self
305    }
306
307    /// Configure the HTTP server.
308    pub fn http_config(mut self, config: crate::http_server::HttpServerConfig) -> Self {
309        self.config.http_health = Some(config);
310        self
311    }
312
313    /// Set the bind address for the HTTP server.
314    pub fn http_bind(mut self, addr: impl Into<SocketAddr>) -> Self {
315        if let Some(http_config) = &mut self.config.http_health {
316            http_config.bind_address = addr.into();
317        } else {
318            self.config.http_health = Some(HttpHealthConfig::new(addr.into()));
319        }
320        self
321    }
322
323    /// Enable HTTP health server with default configuration (port 8081)
324    pub fn http_health(self) -> Self {
325        self.http()
326    }
327
328    /// Configure HTTP health server
329    pub fn http_health_config(self, config: HttpHealthConfig) -> Self {
330        self.http_config(config)
331    }
332
333    /// Set the bind address for HTTP health server
334    pub fn health_bind(self, addr: impl Into<SocketAddr>) -> Self {
335        self.http_bind(addr)
336    }
337
338    pub async fn start(self) -> Result<()> {
339        let (view_index, materialized_registry) =
340            Self::build_view_index_and_registry(self.views, self.materialized_views, &self.spec);
341
342        #[cfg(feature = "otel")]
343        let mut runtime = Runtime::new(self.config, view_index, self.metrics);
344        #[cfg(not(feature = "otel"))]
345        let mut runtime = Runtime::new(self.config, view_index);
346
347        if let Some(plugin) = self.websocket_auth_plugin {
348            runtime = runtime.with_websocket_auth_plugin(plugin);
349        }
350
351        if let Some(plugin) = self.http_auth_plugin {
352            runtime = runtime.with_http_auth_plugin(plugin);
353        }
354
355        if let Some(emitter) = self.websocket_usage_emitter {
356            runtime = runtime.with_websocket_usage_emitter(emitter);
357        }
358
359        if let Some(max_clients) = self.websocket_max_clients {
360            runtime = runtime.with_websocket_max_clients(max_clients);
361        }
362
363        if let Some(rate_limit_config) = self.websocket_rate_limit_config {
364            runtime = runtime.with_websocket_rate_limit_config(rate_limit_config);
365        }
366
367        if let Some(registry) = materialized_registry {
368            runtime = runtime.with_materialized_views(registry);
369        }
370
371        if let Some(spec) = self.spec {
372            runtime = runtime.with_spec(spec);
373        }
374
375        runtime.run().await
376    }
377
378    fn build_view_index_and_registry(
379        views: Option<ViewIndex>,
380        materialized_views: Option<MaterializedViewRegistry>,
381        spec: &Option<Spec>,
382    ) -> (ViewIndex, Option<MaterializedViewRegistry>) {
383        let mut index = views.unwrap_or_default();
384        let mut registry = materialized_views;
385
386        if let Some(ref spec) = spec {
387            let entity_wire_formats = spec
388                .entity_specs
389                .iter()
390                .map(|entity_spec| {
391                    (
392                        entity_spec.state_name.clone(),
393                        ViewSpec::wire_format_from_entity_spec(entity_spec),
394                    )
395                })
396                .collect::<std::collections::HashMap<_, _>>();
397
398            for entity_name in spec.bytecode.entities.keys() {
399                let wire_format = entity_wire_formats
400                    .get(entity_name)
401                    .cloned()
402                    .unwrap_or_default();
403                index.add_spec(ViewSpec {
404                    id: format!("{}/list", entity_name),
405                    export: entity_name.clone(),
406                    mode: Mode::List,
407                    wire_format: wire_format.clone(),
408                    projection: Projection::all(),
409                    filters: Filters::all(),
410                    delivery: Delivery::default(),
411                    pipeline: None,
412                    source_view: None,
413                });
414
415                index.add_spec(ViewSpec {
416                    id: format!("{}/state", entity_name),
417                    export: entity_name.clone(),
418                    mode: Mode::State,
419                    wire_format: wire_format.clone(),
420                    projection: Projection::all(),
421                    filters: Filters::all(),
422                    delivery: Delivery::default(),
423                    pipeline: None,
424                    source_view: None,
425                });
426
427                index.add_spec(ViewSpec {
428                    id: format!("{}/append", entity_name),
429                    export: entity_name.clone(),
430                    mode: Mode::Append,
431                    wire_format,
432                    projection: Projection::all(),
433                    filters: Filters::all(),
434                    delivery: Delivery::default(),
435                    pipeline: None,
436                    source_view: None,
437                });
438            }
439
440            if !spec.views.is_empty() {
441                let reg = registry.get_or_insert_with(MaterializedViewRegistry::new);
442
443                for view_def in &spec.views {
444                    let export = match &view_def.source {
445                        arete_interpreter::ast::ViewSource::Entity { name } => name.clone(),
446                        arete_interpreter::ast::ViewSource::View { id } => {
447                            id.split('/').next().unwrap_or(id).to_string()
448                        }
449                    };
450
451                    let wire_format = entity_wire_formats
452                        .get(&export)
453                        .cloned()
454                        .unwrap_or_default();
455                    let view_spec = ViewSpec::from_view_def(view_def, &export, wire_format);
456                    let pipeline = view_spec.pipeline.clone().unwrap_or_default();
457                    let source_id = view_spec.source_view.clone().unwrap_or_default();
458                    tracing::debug!(
459                        view_id = %view_def.id,
460                        source = %source_id,
461                        "Registering derived view"
462                    );
463
464                    index.add_spec(view_spec);
465
466                    let materialized =
467                        MaterializedView::new(view_def.id.clone(), source_id, pipeline);
468                    reg.register(materialized);
469                }
470            }
471        }
472
473        (index, registry)
474    }
475
476    pub fn build(self) -> Result<Runtime> {
477        let (view_index, materialized_registry) =
478            Self::build_view_index_and_registry(self.views, self.materialized_views, &self.spec);
479
480        #[cfg(feature = "otel")]
481        let mut runtime = Runtime::new(self.config, view_index, self.metrics);
482        #[cfg(not(feature = "otel"))]
483        let mut runtime = Runtime::new(self.config, view_index);
484
485        if let Some(plugin) = self.websocket_auth_plugin {
486            runtime = runtime.with_websocket_auth_plugin(plugin);
487        }
488
489        if let Some(plugin) = self.http_auth_plugin {
490            runtime = runtime.with_http_auth_plugin(plugin);
491        }
492
493        if let Some(max_clients) = self.websocket_max_clients {
494            runtime = runtime.with_websocket_max_clients(max_clients);
495        }
496
497        if let Some(registry) = materialized_registry {
498            runtime = runtime.with_materialized_views(registry);
499        }
500
501        if let Some(spec) = self.spec {
502            runtime = runtime.with_spec(spec);
503        }
504        Ok(runtime)
505    }
506}
507
508#[cfg(test)]
509mod tests {
510    use super::*;
511
512    #[test]
513    fn test_builder_pattern() {
514        let _builder = Server::builder()
515            .websocket()
516            .bind("[::]:8877".parse::<SocketAddr>().unwrap());
517    }
518
519    #[test]
520    fn test_spec_creation() {
521        let bytecode = arete_interpreter::compiler::MultiEntityBytecode::new().build();
522        let spec = Spec::new(bytecode, "test_program");
523        assert_eq!(
524            spec.program_ids.first().map(String::as_str),
525            Some("test_program")
526        );
527    }
528}