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