Explore if we should move the Controller to Rust

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.

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

@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

Feature parity with Python Controller (both Flink and Arroyo support)
Dual interface: Library API + CLI binary
Semantic equivalence: Output configs must be functionally identical to Python
Well-typed configs: Rust structs with serde for validation
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/:

Shared test inputs (JSON/YAML)
Run Python Controller, capture output
Run Rust Controller, capture output
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

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

Explore if we should move the Controller to Rust #101

Description

Activity

milindsrivastava1997 commented on Mar 4, 2026