Explore if we should move the Controller to Rust

[milindsrivastava1997]

Reasoning:

• There is shared query parsing and matching logic between the QueryEngine and the Controller. This logic will keep on expanding and it is best to keep this in the same language
• At some point, the Controller might become an online module that observes incoming queries and data to tweak query plans, and this might need to be higher performance.

[milindsrivastava1997]

Originally by @linear[bot] on 2026-01-08:

PRO-355

[milindsrivastava1997]

Originally by @milindsrivastava1997 on 2026-01-28:

Controller Rust Port Plan

Overview

Port the Python Controller to Rust, providing both a standalone CLI binary and a library for use by QueryEngineRust and other Rust services.

Current State: Python Controller (~1,264 LOC) in Controller/
Target: Rust crate controller_rust in ControllerRust/

Goals

1. Feature parity with Python Controller (both Flink and Arroyo support)
2. Dual interface: Library API + CLI binary
3. Semantic equivalence: Output configs must be functionally identical to Python
4. Well-typed configs: Rust structs with serde for validation
5. Configurable punting: Optional query rejection for performance

---

Phase 0: Cargo Workspace Setup

Create a workspace at the repo root to manage related Rust crates together.

Files to Create/Modify

Create asap-internal/Cargo.toml (workspace root):

[workspace]
resolver = "2"
members = [
    "ControllerRust",
    "QueryEngineRust",
    "CommonDependencies/dependencies/rs/promql_utilities",
    "CommonDependencies/dependencies/rs/sql_utilities",
]

[workspace.package]
edition = "2021"
version = "0.1.0"

[workspace.dependencies]
# Shared dependencies with consistent versions
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
serde_yaml = "0.9"
thiserror = "1.0"
anyhow = "1.0"
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
clap = { version = "4.0", features = ["derive"] }
chrono = { version = "0.4", features = ["serde"] }
indexmap = { version = "2.0", features = ["serde"] }
promql-parser = "0.5.0"

# Internal crates
promql_utilities = { path = "CommonDependencies/dependencies/rs/promql_utilities" }
sql_utilities = { path = "CommonDependencies/dependencies/rs/sql_utilities" }
controller_rust = { path = "ControllerRust" }

Update existing Cargo.toml files to use workspace dependencies:

• QueryEngineRust/Cargo.toml
• CommonDependencies/dependencies/rs/promql_utilities/Cargo.toml
• CommonDependencies/dependencies/rs/sql_utilities/Cargo.toml

Update Utilities/components.conf:

ControllerRust

---

Phase 1: Crate Structure

Directory Layout

ControllerRust/
├── Cargo.toml
├── src/
│   ├── lib.rs                 # Library entry point, public API
│   ├── main.rs                # CLI binary entry point
│   │
│   ├── config/                # Configuration types
│   │   ├── mod.rs
│   │   ├── input.rs           # Input config structs (controller-config.yaml)
│   │   ├── streaming.rs       # StreamingConfig output structs
│   │   ├── inference.rs       # InferenceConfig output structs
│   │   └── sketch_params.rs   # Sketch parameter types
│   │
│   ├── query/                 # Query analysis
│   │   ├── mod.rs
│   │   ├── patterns.rs        # Query pattern definitions
│   │   ├── classification.rs  # Query classification logic
│   │   └── single_query.rs    # SingleQueryConfig equivalent
│   │
│   ├── sketch/                # Sketch parameter logic
│   │   ├── mod.rs
│   │   ├── parameters.rs      # Parameter calculation
│   │   └── window.rs          # Window type selection
│   │
│   ├── engine/                # Streaming engine specifics
│   │   ├── mod.rs
│   │   ├── arroyo.rs          # Arroyo-specific config generation
│   │   └── flink.rs           # Flink-specific config generation
│   │
│   └── error.rs               # Error types
│
├── tests/
│   └── integration.rs         # Integration tests
│
└── README.md

Cargo.toml

[package]
name = "controller_rust"
version.workspace = true
edition.workspace = true

[lib]
name = "controller_rust"
path = "src/lib.rs"

[[bin]]
name = "controller"
path = "src/main.rs"

[dependencies]
# Workspace dependencies
promql_utilities.workspace = true
serde.workspace = true
serde_json.workspace = true
serde_yaml.workspace = true
thiserror.workspace = true
anyhow.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
clap.workspace = true
chrono.workspace = true
indexmap.workspace = true
promql-parser.workspace = true

[dev-dependencies]
tempfile = "3.20"
pretty_assertions = "1.4"

---

Phase 2: Config Types (with Serde Validation)

Input Config (config/input.rs)

use indexmap::IndexMap;
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ControllerConfig {
    pub query_groups: Vec<QueryGroup>,
    pub metrics: Vec<MetricDefinition>,
    #[serde(default)]
    pub sketch_parameters: Option<SketchParameterOverrides>,
    #[serde(default)]
    pub aggregate_cleanup: Option<AggregateCleanupConfig>,
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct QueryGroup {
    pub id: u32,
    pub queries: Vec<String>,
    pub repetition_delay: u64,  // seconds
    #[serde(default)]
    pub client_options: Option<ClientOptions>,
    #[serde(default)]
    pub controller_options: Option<ControllerOptions>,
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ControllerOptions {
    pub accuracy_sla: f64,
    pub latency_sla: f64,
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct MetricDefinition {
    pub metric: String,
    pub labels: Vec<String>,
    #[serde(default)]
    pub cardinality: Option<IndexMap<String, u64>>,
}

// ... additional structs

Streaming Config Output (config/streaming.rs)

use indexmap::IndexMap;
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct StreamingConfig {
    pub aggregations: Vec<AggregationConfig>,
    pub metrics: IndexMap<String, Vec<String>>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct AggregationConfig {
    pub aggregation_id: u32,
    pub aggregation_type: AggregationType,
    pub aggregation_sub_type: String,
    pub metric: String,
    pub spatial_filter: String,
    pub window_type: WindowType,
    pub window_size: u64,
    pub slide_interval: u64,
    pub tumbling_window_size: u64,
    pub parameters: IndexMap<String, serde_json::Value>,
    pub labels: LabelConfig,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LabelConfig {
    pub rollup: Vec<String>,
    pub grouping: Vec<String>,
    pub aggregated: Vec<String>,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub enum AggregationType {
    DeltaSetAggregator,
    CountMinSketch,
    CountMinSketchWithHeap,
    DatasketchesKLL,
    HydraKLL,
    MultipleIncrease,
    Sum,
    Count,
    MinMax,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum WindowType {
    Tumbling,
    Sliding,
}

Inference Config Output (config/inference.rs)

use indexmap::IndexMap;
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InferenceConfig {
    pub queries: Vec<QueryConfig>,
    pub metrics: IndexMap<String, Vec<String>>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryConfig {
    pub query: String,
    pub aggregations: Vec<QueryAggregationRef>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryAggregationRef {
    pub aggregation_id: u32,
    pub num_aggregates_to_retain: Option<u64>,
}

---

Phase 3: Library API Design

Public API (lib.rs)

pub mod config;
pub mod error;

use config::{ControllerConfig, StreamingConfig, InferenceConfig};
use error::ControllerError;

/// Main controller that generates streaming and inference configs
pub struct Controller {
    config: ControllerConfig,
    options: ControllerOptions,
}

#[derive(Debug, Clone, Default)]
pub struct ControllerOptions {
    /// Prometheus scrape interval in seconds
    pub prometheus_scrape_interval: u64,
    /// Target streaming engine
    pub streaming_engine: StreamingEngine,
    /// Enable punting (reject queries that won't perform well)
    pub enable_punting: bool,
}

#[derive(Debug, Clone, Copy, Default)]
pub enum StreamingEngine {
    #[default]
    Arroyo,
    Flink,
}

impl Controller {
    /// Create controller from parsed config
    pub fn new(config: ControllerConfig, options: ControllerOptions) -> Self;

    /// Create controller from YAML file path
    pub fn from_file(path: &Path, options: ControllerOptions) -> Result<Self, ControllerError>;

    /// Create controller from YAML string
    pub fn from_yaml(yaml: &str, options: ControllerOptions) -> Result<Self, ControllerError>;

    /// Generate both configs
    pub fn generate(&self) -> Result<ControllerOutput, ControllerError>;

    /// Generate and write to directory
    pub fn generate_to_dir(&self, output_dir: &Path) -> Result<ControllerOutput, ControllerError>;
}

#[derive(Debug, Clone)]
pub struct ControllerOutput {
    pub streaming_config: StreamingConfig,
    pub inference_config: InferenceConfig,
    /// Queries that were punted (if punting enabled)
    pub punted_queries: Vec<PuntedQuery>,
}

#[derive(Debug, Clone)]
pub struct PuntedQuery {
    pub query: String,
    pub reason: String,
}

Usage from QueryEngineRust

use controller_rust::{Controller, ControllerOptions, StreamingEngine};

// Option 1: From file
let controller = Controller::from_file(
    Path::new("controller-config.yaml"),
    ControllerOptions {
        prometheus_scrape_interval: 15,
        streaming_engine: StreamingEngine::Arroyo,
        enable_punting: true,
    },
)?;
let output = controller.generate()?;

// Option 2: From parsed config (if QueryEngine already has it)
let controller = Controller::new(config, options);
let output = controller.generate()?;

// Access the configs
let streaming = output.streaming_config;
let inference = output.inference_config;

---

Phase 4: CLI Design

Command Line Interface

controller-rust 0.1.0
ASAP Controller - generates streaming and inference configs from query workloads

USAGE:
    controller [OPTIONS] --input <FILE>

OPTIONS:
    -i, --input <FILE>              Input controller config YAML file
    -o, --output-dir <DIR>          Output directory [default: .]
    -s, --scrape-interval <SECS>    Prometheus scrape interval [default: 15]
    -e, --engine <ENGINE>           Streaming engine: arroyo|flink [default: arroyo]
        --enable-punting            Enable query punting for performance
        --streaming-only            Only output streaming_config.yaml
        --inference-only            Only output inference_config.yaml
        --json                      Output as JSON instead of YAML
    -v, --verbose                   Enable verbose logging
    -h, --help                      Print help
    -V, --version                   Print version

Implementation (main.rs)

use clap::Parser;
use controller_rust::{Controller, ControllerOptions, StreamingEngine};
use std::path::PathBuf;

#[derive(Parser)]
#[command(name = "controller")]
#[command(about = "ASAP Controller - generates streaming and inference configs")]
struct Cli {
    #[arg(short, long, value_name = "FILE")]
    input: PathBuf,

    #[arg(short, long, value_name = "DIR", default_value = ".")]
    output_dir: PathBuf,

    #[arg(short, long, default_value = "15")]
    scrape_interval: u64,

    #[arg(short, long, default_value = "arroyo")]
    engine: StreamingEngine,

    #[arg(long)]
    enable_punting: bool,

    #[arg(long)]
    json: bool,

    #[arg(short, long)]
    verbose: bool,
}

fn main() -> anyhow::Result<()> {
    let cli = Cli::parse();

    // Setup tracing
    if cli.verbose {
        tracing_subscriber::fmt()
            .with_env_filter("controller_rust=debug")
            .init();
    }

    let options = ControllerOptions {
        prometheus_scrape_interval: cli.scrape_interval,
        streaming_engine: cli.engine,
        enable_punting: cli.enable_punting,
    };

    let controller = Controller::from_file(&cli.input, options)?;
    let output = controller.generate_to_dir(&cli.output_dir)?;

    // Report results
    println!("Generated {} aggregations", output.streaming_config.aggregations.len());
    println!("Generated {} query configs", output.inference_config.queries.len());

    if !output.punted_queries.is_empty() {
        println!("Punted {} queries:", output.punted_queries.len());
        for pq in &output.punted_queries {
            println!("  - {}: {}", pq.query, pq.reason);
        }
    }

    Ok(())
}

---

Phase 5: Testing Infrastructure

Test Strategy

Use the same pattern as CommonDependencies/tests/compare_matched_tokens/:

1. Shared test inputs (JSON/YAML)
2. Run Python Controller, capture output
3. Run Rust Controller, capture output
4. Compare outputs semantically

Directory Structure

ControllerRust/
└── tests/
    ├── comparison/
    │   ├── test_data/
    │   │   ├── test_cases.json          # Test case definitions
    │   │   └── configs/                  # Input configs from experiment_outputs
    │   │       ├── case_001.yaml
    │   │       ├── case_002.yaml
    │   │       └── ...
    │   ├── python_runner.py              # Runs Python controller
    │   ├── rust_runner.rs                # Runs Rust controller (integration test)
    │   ├── comparator.py                 # Compares outputs semantically
    │   └── master_test_runner.py         # Orchestrates everything
    └── integration.rs                    # Rust-only integration tests

Test Case Format (test_cases.json)

{
  "test_cases": [
    {
      "id": "final_generated_workload_1",
      "description": "Mixed spatial and temporal queries",
      "input_config": "configs/case_001.yaml",
      "prometheus_scrape_interval": 15,
      "streaming_engine": "arroyo",
      "enable_punting": false
    },
    {
      "id": "quantile_queries",
      "description": "Quantile-heavy workload",
      "input_config": "configs/case_002.yaml",
      "prometheus_scrape_interval": 1,
      "streaming_engine": "arroyo",
      "enable_punting": true
    }
  ]
}

Semantic Comparison (comparator.py)

def compare_streaming_configs(python_config, rust_config) -> ComparisonResult:
    """Compare streaming configs semantically, not byte-for-byte."""

    issues = []

    # Compare aggregation count
    if len(python_config['aggregations']) != len(rust_config['aggregations']):
        issues.append(f"Aggregation count mismatch: {len(python_config['aggregations'])} vs {len(rust_config['aggregations'])}")

    # Build lookup by (metric, aggregationType, labels hash)
    py_aggs = {agg_key(a): a for a in python_config['aggregations']}
    rs_aggs = {agg_key(a): a for a in rust_config['aggregations']}

    # Compare each aggregation
    for key in py_aggs:
        if key not in rs_aggs:
            issues.append(f"Missing aggregation in Rust: {key}")
            continue

        py_agg, rs_agg = py_aggs[key], rs_aggs[key]

        # Compare numeric fields
        for field in ['windowSize', 'slideInterval', 'tumblingWindowSize']:
            if py_agg[field] != rs_agg[field]:
                issues.append(f"{key}: {field} mismatch: {py_agg[field]} vs {rs_agg[field]}")

        # Compare parameters
        if py_agg['parameters'] != rs_agg['parameters']:
            issues.append(f"{key}: parameters mismatch")

        # Compare labels (as sets, order doesn't matter)
        for label_type in ['rollup', 'grouping', 'aggregated']:
            py_labels = set(py_agg['labels'][label_type])
            rs_labels = set(rs_agg['labels'][label_type])
            if py_labels != rs_labels:
                issues.append(f"{key}: {label_type} labels mismatch")

    return ComparisonResult(
        passed=len(issues) == 0,
        issues=issues
    )

Test Corpus Selection

Extract test configs from experiment_outputs/:

Source Directory	Test Focus
final_generated_workload_prometheus_*	Mixed workloads
qot_*	quantile_over_time patterns
quantile_1s_*	Spatial quantile patterns
non_quantile_*	Sum/count/rate patterns
diverse_test_*	Edge cases

Script to extract configs:

#!/bin/bash
# Extract controller input configs for testing
for dir in experiment_outputs/final_generated_workload_prometheus_20251204_231532_{1..5}; do
    name=$(basename $dir)
    cp "$dir/controller_client_configs/sketchdb.yaml" \
       "ControllerRust/tests/comparison/test_data/configs/${name}.yaml"
done

---

Phase 6: Implementation Order

Step 1: Workspace Setup

• Create asap-internal/Cargo.toml workspace file
• Update QueryEngineRust/Cargo.toml to use workspace deps
• Update promql_utilities/Cargo.toml to use workspace deps
• Update sql_utilities/Cargo.toml to use workspace deps
• Verify cargo build works from repo root

Step 2: Scaffold ControllerRust

• Create ControllerRust/ directory structure
• Create Cargo.toml with workspace deps
• Implement config structs (config/ module)
• Implement error types (error.rs)
• Create stub lib.rs and main.rs

Step 3: Core Logic

• Port query pattern definitions (query/patterns.rs)
• Port query classification (query/classification.rs)
• Port SingleQueryConfig logic (query/single_query.rs)
• Port sketch parameter calculation (sketch/parameters.rs)
• Port window logic (sketch/window.rs)

Step 4: Config Generation

• Implement Controller::generate() orchestration
• Implement streaming config generation
• Implement inference config generation
• Implement aggregation deduplication
• Implement punting logic (configurable)

Step 5: CLI

• Implement CLI argument parsing
• Implement file I/O (YAML read/write)
• Add tracing/logging
• Test manually against Python output

Step 6: Testing Infrastructure

• Extract test configs from experiment_outputs
• Create test_cases.json
• Implement python_runner.py
• Implement comparator.py
• Implement master_test_runner.py
• Run comparison tests, fix discrepancies

Step 7: Integration

• Add controller_rust dependency to QueryEngineRust/Cargo.toml
• Update Utilities/components.conf
• Create Dockerfile for ControllerRust
• Update any scripts that invoke Python Controller

---

Dependencies Summary

From Workspace

• promql_utilities - PromQL parsing, pattern matching, query classification
• serde, serde_yaml, serde_json - Serialization
• thiserror, anyhow - Error handling
• tracing, tracing-subscriber - Logging
• clap - CLI parsing
• chrono - Time handling
• indexmap - Ordered maps for deterministic YAML output

New Dependencies

• promql-parser = "0.5.0" - Already used by promql_utilities

---

Risks and Mitigations

Risk	Mitigation
Pattern matching behavior differs	Comprehensive comparison tests against Python
YAML output format differences	Use indexmap for key ordering; comparison tests ignore cosmetic differences
Edge cases in query classification	Extract diverse test cases from experiment_outputs
promql_utilities API gaps	Review Python usage, add missing Rust APIs if needed

---

Estimated LOC

Module	Estimated LOC
config/ (all structs)	300-400
query/ (patterns, classification)	400-500
sketch/ (parameters, windows)	200-300
engine/ (arroyo, flink)	100-150
lib.rs + error.rs	150-200
main.rs (CLI)	100-150
Total	1,250-1,700 LOC

Plus ~300-400 LOC for test infrastructure (Python comparator, test runner).

---

Success Criteria

1. cargo build succeeds for workspace
2. cargo test passes all unit tests
3. CLI produces output for all test configs
4. Comparison tests show semantic equivalence with Python for 100% of test cases
5. QueryEngineRust can import and use controller_rust as a library