Feature 3: Create Storage Controller Client Library

[bittermandel]

Purpose

Create a robust HTTP client library for storage controller API integration with proper retry logic, concurrency controls, and error handling.
Inspired by https://github.com/neondatabase/neon/blob/main/control_plane/storcon_cli/src/main.rs

Success Criteria

• All API methods (drain, fill, status) have proper retry logic
• JWT token refresh works correctly
• Error categorization properly distinguishes retryable vs fatal errors

[linear]

NEON-26 Feature 3: Create Storage Controller Client Library

Purpose

Create a robust HTTP client library for storage controller API integration with proper retry logic, concurrency controls, and error handling.

Required Result

• New crate `crates/storage_controller_client/` with HTTP client
• Exponential backoff retry strategy for API calls
• Concurrency limits via semaphores to prevent storage controller overload
• JWT token management with refresh capabilities
• Proper error categorization (retryable vs fatal)
• Comprehensive error handling for all storage controller APIs

Implementation Details

1. Create New Crate

Create `crates/storage_controller_client/Cargo.toml`:

```toml
[package]
name = "storage_controller_client"
version = "0.1.0"
edition = "2021"

[dependencies]
reqwest = { version = "0.11", features = ["json"] }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
tokio = { version = "1.0", features = ["full"] }
tokio-retry = "0.3"
tracing = "0.1"
thiserror = "1.0"
```

2. Core Client Implementation

Create `crates/storage_controller_client/src/lib.rs`:

```rust
use reqwest::Client as HttpClient;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{RwLock, Semaphore};
use tokio_retry::strategy::{ExponentialBackoff, jitter};
use tokio_retry::Retry;
use tracing::{debug, warn, error};

#[derive(Debug, thiserror::Error)]
pub enum StorageControllerError {
#[error("Retryable error: {0}")]
Retryable(String),

#[error("Fatal error: {0}")]
Fatal(String),

#[error("HTTP error: {0}")]
Http(#[from] reqwest::Error),

#[error("Serialization error: {0}")]
Serialization(#[from] serde_json::Error),

#[error("Semaphore error: {0}")]
Semaphore(#[from] tokio::sync::AcquireError),

}

pub type Result = std::result::Result;

#[derive(Debug, Clone)] pub struct StorageControllerClient { base_url: String, jwt_token: Arc>, http_client: HttpClient, // Semaphores to limit concurrent operations drain_semaphore: Arc, fill_semaphore: Arc, general_semaphore: Arc, }

#[derive(Debug, Serialize, Deserialize)] pub struct NodeStatus { pub scheduling_policy: String, pub drain_status: Option, pub fill_status: Option, }

#[derive(Debug, Serialize, Deserialize)] pub struct DrainStatus { pub in_progress: bool, pub started_at: Option, pub remaining_shards: u32, pub total_shards: u32, }

#[derive(Debug, Serialize, Deserialize)] pub struct FillStatus { pub in_progress: bool, pub started_at: Option, pub received_shards: u32, pub target_shards: u32, }

impl StorageControllerClient {
pub fn new(base_url: String, jwt_token: String, max_concurrent_ops: usize) -> Self {
Self {
base_url,
jwt_token: Arc::new(RwLock::new(jwt_token)),
http_client: HttpClient::new(),
drain_semaphore: Arc::new(Semaphore::new(max_concurrent_ops)),
fill_semaphore: Arc::new(Semaphore::new(max_concurrent_ops)),
general_semaphore: Arc::new(Semaphore::new(max_concurrent_ops * 2)),
}
}

/// Start draining a pageserver node
pub async fn start_drain(&self, node_id: &str) -> Result<()> {
    let _permit = self.drain_semaphore.acquire().await?;
    
    let retry_strategy = ExponentialBackoff::from_millis(100)
        .max_delay(Duration::from_secs(30))
        .take(5)
        .map(jitter);

    Retry::spawn(retry_strategy, || async {
        self.start_drain_inner(node_id).await
    }).await
}

async fn start_drain_inner(&self, node_id: &str) -> Result<()> {
    let url = format!("{}/control/v1/node/{}/drain", self.base_url, node_id);
    let token = self.jwt_token.read().await.clone();
    
    debug!("Starting drain for node {}", node_id);
    
    let response = self.http_client
        .put(&url)
        .bearer_auth(&token)
        .send()
        .await?;
        
    match response.status() {
        reqwest::StatusCode::OK => {
            debug!("Drain started successfully for node {}", node_id);
            Ok(())
        }
        reqwest::StatusCode::CONFLICT => {
            // Already draining - idempotent success
            debug!("Node {} already draining", node_id);
            Ok(())
        }
        reqwest::StatusCode::TOO_MANY_REQUESTS => {
            warn!("Rate limited when starting drain for node {}", node_id);
            Err(StorageControllerError::Retryable("Rate limited".to_string()))
        }
        status if status.is_server_error() => {
            warn!("Server error when starting drain for node {}: {}", node_id, status);
            Err(StorageControllerError::Retryable(format!("Server error: {}", status)))
        }
        status => {
            error!("Failed to start drain for node {}: {}", node_id, status);
            Err(StorageControllerError::Fatal(format!("Failed to start drain: {}", status)))
        }
    }
}

/// Start filling a pageserver node
pub async fn start_fill(&self, node_id: &str) -> Result<()> {
    let _permit = self.fill_semaphore.acquire().await?;
    
    let retry_strategy = ExponentialBackoff::from_millis(100)
        .max_delay(Duration::from_secs(30))
        .take(5)
        .map(jitter);

    Retry::spawn(retry_strategy, || async {
        self.start_fill_inner(node_id).await
    }).await
}

async fn start_fill_inner(&self, node_id: &str) -> Result<()> {
    let url = format!("{}/control/v1/node/{}/fill", self.base_url, node_id);
    let token = self.jwt_token.read().await.clone();
    
    debug!("Starting fill for node {}", node_id);
    
    let response = self.http_client
        .put(&url)
        .bearer_auth(&token)
        .send()
        .await?;
        
    match response.status() {
        reqwest::StatusCode::OK => {
            debug!("Fill started successfully for node {}", node_id);
            Ok(())
        }
        reqwest::StatusCode::CONFLICT => {
            // Already filling - idempotent success
            debug!("Node {} already filling", node_id);
            Ok(())
        }
        reqwest::StatusCode::TOO_MANY_REQUESTS => {
            warn!("Rate limited when starting fill for node {}", node_id);
            Err(StorageControllerError::Retryable("Rate limited".to_string()))
        }
        status if status.is_server_error() => {
            warn!("Server error when starting fill for node {}: {}", node_id, status);
            Err(StorageControllerError::Retryable(format!("Server error: {}", status)))
        }
        status => {
            error!("Failed to start fill for node {}: {}", node_id, status);
            Err(StorageControllerError::Fatal(format!("Failed to start fill: {}", status)))
        }
    }
}

/// Get node status
pub async fn get_node_status(&self, node_id: &str) -> Result<NodeStatus> {
    let _permit = self.general_semaphore.acquire().await?;
    
    let retry_strategy = ExponentialBackoff::from_millis(100)
        .max_delay(Duration::from_secs(10))
        .take(3)
        .map(jitter);

    Retry::spawn(retry_strategy, || async {
        self.get_node_status_inner(node_id).await
    }).await
}

async fn get_node_status_inner(&self, node_id: &str) -> Result<NodeStatus> {
    let url = format!("{}/control/v1/node/{}", self.base_url, node_id);
    let token = self.jwt_token.read().await.clone();
    
    let response = self.http_client
        .get(&url)
        .bearer_auth(&token)
        .send()
        .await?;
        
    match response.status() {
        reqwest::StatusCode::OK => {
            let status: NodeStatus = response.json().await?;
            Ok(status)
        }
        reqwest::StatusCode::TOO_MANY_REQUESTS => {
            Err(StorageControllerError::Retryable("Rate limited".to_string()))
        }
        status if status.is_server_error() => {
            Err(StorageControllerError::Retryable(format!("Server error: {}", status)))
        }
        status => {
            Err(StorageControllerError::Fatal(format!("Failed to get node status: {}", status)))
        }
    }
}

/// Refresh JWT token
pub async fn refresh_token(&self, new_token: String) {
    let mut token = self.jwt_token.write().await;
    *token = new_token;
    debug!("JWT token refreshed");
}

/// Check if drain is complete
pub async fn is_drain_complete(&self, node_id: &str) -> Result<bool> {
    let status = self.get_node_status(node_id).await?;
    Ok(status.scheduling_policy == "PauseForRestart" && 
       status.drain_status.map_or(0, |d| d.remaining_shards) == 0)
}

/// Check if fill is complete
pub async fn is_fill_complete(&self, node_id: &str) -> Result<bool> {
    let status = self.get_node_status(node_id).await?;
    Ok(status.scheduling_policy == "Active" && 
       status.fill_status.map_or(true, |f| !f.in_progress))
}

}

#[cfg(test)]
mod tests {
use super::*;

#[tokio::test]
async fn test_client_creation() {
    let client = StorageControllerClient::new(
        "http://localhost:8080".to_string(),
        "test-token".to_string(),
        10
    );
    
    assert_eq!(client.base_url, "http://localhost:8080");
    let token = client.jwt_token.read().await;
    assert_eq!(*token, "test-token");
}

#[tokio::test]
async fn test_token_refresh() {
    let client = StorageControllerClient::new(
        "http://localhost:8080".to_string(),
        "old-token".to_string(),
        10
    );
    
    client.refresh_token("new-token".to_string()).await;
    
    let token = client.jwt_token.read().await;
    assert_eq!(*token, "new-token");
}

}
```

Testing Requirements

Unit Tests

• Test client creation and configuration
• Test JWT token refresh
• Test error categorization logic
• Test retry strategy behavior

Mock Server Tests

Create `crates/storage_controller_client/tests/integration_tests.rs`:

```rust
use storage_controller_client::{StorageControllerClient, NodeStatus};
use std::time::Duration;
use tokio::time::timeout;

// Mock server setup helper
async fn setup_mock_server() -> (String, tokio::task::JoinHandle<()>) {
// Implementation would start a mock HTTP server
// Return (base_url, server_handle)
todo!("Implement mock server for testing")
}

#[tokio::test] async fn test_drain_api_success() -> Result<(), Box> { let (base_url, _server) = setup_mock_server().await; let client = StorageControllerClient::new(base_url, "test-token".to_string(), 5);

// Test successful drain
let result = timeout(Duration::from_secs(10), client.start_drain("node-123")).await?;
assert!(result.is_ok());

Ok(())

}

#[tokio::test] async fn test_retry_on_server_error() -> Result<(), Box> { let (base_url, _server) = setup_mock_server().await; let client = StorageControllerClient::new(base_url, "test-token".to_string(), 5);

// Mock server should return 500 then 200
let result = timeout(Duration::from_secs(30), client.start_drain("node-retry")).await?;
assert!(result.is_ok());

Ok(())

}

#[tokio::test] async fn test_concurrency_limits() -> Result<(), Box> { let (base_url, _server) = setup_mock_server().await; let client = StorageControllerClient::new(base_url, "test-token".to_string(), 2);

// Start 5 concurrent drain operations (should be limited to 2)
let futures: Vec<_> = (0..5).map(|i| {
    let client = client.clone();
    tokio::spawn(async move {
        client.start_drain(&format!("node-{}", i)).await
    })
}).collect();

// All should complete successfully (some after waiting)
for future in futures {
    let result = future.await?;
    assert!(result.is_ok());
}

Ok(())

}
```

E2E Tests

Add to existing E2E test suite:

```rust #[tokio::test] async fn test_storage_controller_client_integration() -> Result<(), Box> { let test_env = TestEnv::new("storage-controller-client").await?;

// This test requires a running storage controller
// Skip if not available in test environment
if !storage_controller_available(&test_env).await {
    return Ok(());
}

let client = StorageControllerClient::new(
    "http://storage-controller:8080".to_string(),
    "test-jwt-token".to_string(),
    5
);

// Test basic connectivity
let status = client.get_node_status("test-node").await;
// Should either succeed or fail with a known error
assert!(status.is_ok() || matches!(status.unwrap_err(), StorageControllerError::Fatal(_)));

test_env.cleanup().await?;
Ok(())

}
```

Success Criteria

• Client can be instantiated with proper configuration
• All API methods (drain, fill, status) have proper retry logic
• Concurrency limits prevent storage controller overload
• JWT token refresh works correctly
• Error categorization properly distinguishes retryable vs fatal errors
• Unit tests cover all error scenarios
• Integration tests with mock server pass
• Library can be used by other crates in the workspace

Files to Create

• `crates/storage_controller_client/Cargo.toml`
• `crates/storage_controller_client/src/lib.rs`
• `crates/storage_controller_client/tests/integration_tests.rs`

Files to Modify

• `Cargo.toml` (add new crate to workspace)

Dependencies

• Independent feature that can be implemented in parallel