forge_orchestration/sdk/
mod.rs

1//! # Forge SDK
2//!
3//! Lightweight SDK for workloads to interact with the Forge orchestrator.
4//!
5//! This module provides utilities for workloads running under Forge orchestration:
6//! - **Lifecycle**: Signal readiness, handle graceful shutdown
7//! - **Port Allocation**: Request and release ports dynamically
8//! - **Client**: HTTP client for Forge API communication
9//!
10//! ## Quick Start
11//!
12//! ```rust,no_run
13//! use forge_orchestration::sdk::{ready, allocate_port, graceful_shutdown};
14//!
15//! #[tokio::main]
16//! async fn main() -> forge_orchestration::Result<()> {
17//!     // Signal readiness
18//!     ready()?;
19//!
20//!     // Get a port
21//!     let port = allocate_port(8000..9000)?;
22//!     println!("Listening on port {}", port);
23//!
24//!     // Handle shutdown
25//!     graceful_shutdown();
26//!     Ok(())
27//! }
28//! ```
29
30pub mod client;
31pub mod lifecycle;
32pub mod port;
33pub mod session;
34pub mod spot;
35
36pub use client::{ForgeClient, MetricsReport, start_heartbeat};
37pub use lifecycle::{graceful_shutdown, is_ready, notify_shutdown, ready, shutdown_requested, shutdown_signal};
38pub use port::{
39    allocate_port, allocate_specific_port, allocated_ports, is_port_available, release_port,
40    allocate_udp_port, allocate_specific_udp_port, allocate_game_port, allocate_specific_game_port,
41    allocate_port_with_protocol, allocate_specific_port_with_protocol,
42    allocated_ports_detailed, is_udp_port_available, is_port_available_for_protocol,
43    AllocatedPort, PortAllocator, Protocol,
44};
45pub use spot::{
46    SpotHandler, SpotInterruption, SpotAction, CloudProvider,
47    start_spot_monitor, wait_for_spot_interruption, is_spot_instance,
48};
49pub use session::{
50    Session, SessionId, SessionState, SessionTracker, SessionConfig,
51    ConnectionInfo,
52};
53
54use crate::error::ForgeError;
55
56/// Environment variable for Forge API endpoint
57pub const FORGE_API_ENV: &str = "FORGE_API";
58
59/// Environment variable for allocation ID
60pub const FORGE_ALLOC_ID_ENV: &str = "FORGE_ALLOC_ID";
61
62/// Environment variable for task name
63pub const FORGE_TASK_NAME_ENV: &str = "FORGE_TASK_NAME";
64
65/// Get the Forge API endpoint from environment
66pub fn forge_api_url() -> Option<String> {
67    std::env::var(FORGE_API_ENV).ok()
68}
69
70/// Get the allocation ID from environment
71pub fn alloc_id() -> Option<String> {
72    std::env::var(FORGE_ALLOC_ID_ENV).ok()
73}
74
75/// Get the task name from environment
76pub fn task_name() -> Option<String> {
77    std::env::var(FORGE_TASK_NAME_ENV).ok()
78}
79
80/// SDK-specific error type
81#[derive(Debug, thiserror::Error)]
82pub enum SdkError {
83    /// Forge API not configured
84    #[error("Forge API not configured: set {0} environment variable")]
85    NotConfigured(&'static str),
86
87    /// API communication error
88    #[error("API error: {0}")]
89    Api(String),
90
91    /// Port allocation failed
92    #[error("Port allocation failed: {0}")]
93    PortAllocation(String),
94
95    /// Lifecycle error
96    #[error("Lifecycle error: {0}")]
97    Lifecycle(String),
98
99    /// IO error
100    #[error("IO error: {0}")]
101    Io(#[from] std::io::Error),
102
103    /// Serialization error
104    #[error("Serialization error: {0}")]
105    Serialization(#[from] serde_json::Error),
106
107    /// HTTP error
108    #[error("HTTP error: {0}")]
109    Http(#[from] reqwest::Error),
110}
111
112impl SdkError {
113    /// Create an API error
114    pub fn api(msg: impl Into<String>) -> Self {
115        Self::Api(msg.into())
116    }
117
118    /// Create a port allocation error
119    pub fn port(msg: impl Into<String>) -> Self {
120        Self::PortAllocation(msg.into())
121    }
122
123    /// Create a lifecycle error
124    pub fn lifecycle(msg: impl Into<String>) -> Self {
125        Self::Lifecycle(msg.into())
126    }
127}
128
129impl From<SdkError> for ForgeError {
130    fn from(err: SdkError) -> Self {
131        ForgeError::Internal(err.to_string())
132    }
133}
134
135/// SDK Result type alias
136pub type SdkResult<T> = std::result::Result<T, SdkError>;
forge_orchestration/sdk/mod.rs

forge_orchestration/sdk/
mod.rs
