Skip to content

Commit 472de6e

Browse files
authored
Merge pull request #525 from ProjectASAP/feat/o11ybench-execution
feat(o11y): execute real inputs through Planner-selected backend plans
2 parents 89ca0f5 + 6179df9 commit 472de6e

11 files changed

Lines changed: 706 additions & 15 deletions

File tree

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
//! Control-plane entry point: cost-select a workload and emit its atomic install request.
2+
use control_plane::physical::compiler::BackendLocalPlanningSnapshot;
3+
use serde_json::json;
4+
5+
fn main() -> Result<(), Box<dyn std::error::Error>> {
6+
let path = std::env::args()
7+
.nth(1)
8+
.ok_or("usage: compile_workload_artifact SNAPSHOT.json")?;
9+
let snapshot: BackendLocalPlanningSnapshot = serde_json::from_slice(&std::fs::read(path)?)?;
10+
if snapshot.snapshot_version != 2 {
11+
return Err(
12+
"execution evaluation requires version 2 complete workload cost evidence".into(),
13+
);
14+
}
15+
let start = std::time::Instant::now();
16+
let plan = snapshot.compile()?;
17+
let elapsed = start.elapsed().as_nanos();
18+
let comparison = plan
19+
.cost_comparison
20+
.ok_or("missing complete-plan comparison")?;
21+
println!(
22+
"{}",
23+
serde_json::to_string_pretty(&json!({
24+
"schema_version": 1,
25+
"planning_elapsed_ns": elapsed,
26+
"envelope": plan.envelope,
27+
"cost_comparison": comparison,
28+
"lifecycle_estimates": plan.lifecycle_estimates,
29+
"install_request": {
30+
"precompute_plan": plan.precompute_plan,
31+
"transmission_plan": plan.transmission_plan,
32+
"backend_plan": plan.backend_plan.encode_to_vec(),
33+
"query_plan": plan.query_plan,
34+
"storage_routing": null,
35+
"adaptation_evidence": []
36+
}
37+
}))?
38+
);
39+
Ok(())
40+
}

‎data_plane/src/drivers/ingest/prometheus_remote_write.rs‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,7 @@ struct ReceiverInner {
100100

101101
#[derive(Default)]
102102
struct DedupState {
103+
input_closed: bool,
103104
values: HashMap<(u64, u64, String, i64), DedupValue>,
104105
expiry: VecDeque<(Instant, u64, u64, String, i64)>,
105106
}
@@ -112,6 +113,8 @@ enum DedupValue {
112113

113114
#[derive(Debug, thiserror::Error)]
114115
pub enum RemoteWriteError {
116+
#[error("finite input has been closed by the drain barrier")]
117+
InputClosed,
115118
#[error("compressed request exceeds {0} bytes")]
116119
CompressedTooLarge(usize),
117120
#[error("decompressed request exceeds {0} bytes")]
@@ -171,6 +174,15 @@ impl PrometheusRemoteWriteReceiver {
171174
self.inner.stats.clone()
172175
}
173176

177+
/// Permanently seal this finite source before queuing worker barriers.
178+
pub async fn drain(&self) -> Result<(), String> {
179+
{
180+
let mut state = self.inner.dedup.lock().map_err(|e| e.to_string())?;
181+
state.input_closed = true;
182+
}
183+
self.inner.ingest.router.drain().await
184+
}
185+
174186
/// Decode, validate, deduplicate, and enqueue one whole v1 request.
175187
/// All validation and all queue reservations complete before any message
176188
/// becomes visible to a worker.
@@ -235,6 +247,9 @@ impl PrometheusRemoteWriteReceiver {
235247
.dedup
236248
.lock()
237249
.expect("remote write dedup poisoned");
250+
if dedup.input_closed {
251+
return Err(RemoteWriteError::InputClosed);
252+
}
238253
dedup.evict_before(now.checked_sub(config.dedup_horizon).unwrap_or(now));
239254

240255
// Validate conflicts both against committed history and inside this
@@ -664,6 +679,28 @@ mod tests {
664679
})
665680
}
666681

682+
// Closing finite input prevents writes racing behind the completion barrier.
683+
#[tokio::test]
684+
async fn finite_input_drain_seals_receiver_and_propagates_worker_failure() {
685+
let (receiver, mut worker) = configured_receiver();
686+
receiver.accept(&one_sample(1.0)).unwrap();
687+
let handle = receiver.clone();
688+
let drain = tokio::spawn(async move { handle.drain().await });
689+
assert!(matches!(
690+
worker.recv().await.unwrap(),
691+
WorkerMessage::GroupSamples { .. }
692+
));
693+
let WorkerMessage::Drain(reply) = worker.recv().await.unwrap() else {
694+
panic!("expected barrier")
695+
};
696+
assert!(matches!(
697+
receiver.accept(&one_sample(1.0)),
698+
Err(RemoteWriteError::InputClosed)
699+
));
700+
reply.send(Err("sink write failed".into())).unwrap();
701+
assert_eq!(drain.await.unwrap().unwrap_err(), "sink write failed");
702+
}
703+
667704
#[test]
668705
fn rejects_writes_without_an_active_physical_plan() {
669706
let (sender, _worker) = mpsc::channel(1);

‎data_plane/src/drivers/query/fallback/mod.rs‎

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ pub enum FallbackResponse {
1919

2020
impl IntoResponse for FallbackResponse {
2121
fn into_response(self) -> Response {
22-
match self {
22+
let mut response = match self {
2323
FallbackResponse::Json(value) => Json(value).into_response(),
2424
FallbackResponse::Text(text) => {
2525
// Return plain text with appropriate content type
@@ -32,7 +32,28 @@ impl IntoResponse for FallbackResponse {
3232
)
3333
.into_response()
3434
}
35-
}
35+
};
36+
response.headers_mut().insert(
37+
"x-asap-execution",
38+
axum::http::HeaderValue::from_static("exact_fallback"),
39+
);
40+
response
41+
}
42+
}
43+
44+
#[cfg(test)]
45+
mod execution_attribution_tests {
46+
use super::*;
47+
48+
// An HTTP fallback must be distinguishable from warm execution, even if
49+
// its upstream response carries misleading or absent source annotations.
50+
#[test]
51+
fn forwarded_response_has_backend_owned_execution_marker() {
52+
let response = FallbackResponse::Json(serde_json::json!({
53+
"status": "success", "infos": ["data_source: asap_query"]
54+
}))
55+
.into_response();
56+
assert_eq!(response.headers()["x-asap-execution"], "exact_fallback");
3657
}
3758
}
3859

‎data_plane/src/drivers/query/servers/http.rs‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -470,6 +470,7 @@ impl HttpServer {
470470
.layer(DefaultBodyLimit::max(request_body_limit)),
471471
)
472472
.route("/api/v1/store/metrics", get(handle_store_metrics))
473+
.route("/api/v1/precompute/drain", post(handle_precompute_drain))
473474
.route(
474475
"/api/v1/streaming-config",
475476
get(handle_get_streaming_config).post(handle_post_streaming_config),
@@ -5577,6 +5578,8 @@ async fn handle_prometheus_remote_write(
55775578

55785579
match receiver.accept(&body) {
55795580
Ok(()) => StatusCode::NO_CONTENT.into_response(),
5581+
Err(crate::drivers::ingest::prometheus_remote_write::RemoteWriteError::InputClosed) =>
5582+
(StatusCode::CONFLICT, "finite input is closed").into_response(),
55805583
Err(
55815584
crate::drivers::ingest::prometheus_remote_write::RemoteWriteError::CompressedTooLarge(
55825585
_,
@@ -5636,6 +5639,29 @@ async fn handle_health(State(state): State<AppState>) -> axum::response::Respons
56365639
(StatusCode::OK, "ok").into_response()
56375640
}
56385641

5642+
/// Explicit end-of-input; this seals Remote Write for the lifetime of the process.
5643+
async fn handle_precompute_drain(State(state): State<AppState>) -> Response {
5644+
let Some(receiver) = state.remote_write.as_ref() else {
5645+
return (StatusCode::NOT_FOUND, "Remote Write is disabled").into_response();
5646+
};
5647+
match receiver.drain().await {
5648+
Ok(()) => (
5649+
StatusCode::OK,
5650+
axum::Json(serde_json::json!({
5651+
"status": "success", "input_closed": true, "complete": true
5652+
})),
5653+
)
5654+
.into_response(),
5655+
Err(error) => (
5656+
StatusCode::INTERNAL_SERVER_ERROR,
5657+
axum::Json(serde_json::json!({
5658+
"status": "error", "input_closed": true, "complete": false, "error": error
5659+
})),
5660+
)
5661+
.into_response(),
5662+
}
5663+
}
5664+
56395665
/// Return list of metrics currently in the store.
56405666
async fn handle_store_metrics(State(state): State<AppState>) -> axum::response::Response {
56415667
use axum::http::StatusCode;

‎data_plane/src/precompute_engine/output_sink.rs‎

Lines changed: 23 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -100,11 +100,9 @@ impl SketchStoreSink {
100100
}
101101
}
102102

103-
/// Best-effort write to `SketchStore` for one PrecomputedOutput.
104-
/// Logs and skips on missing agg_config or other transient
105-
/// inconsistencies — a SketchStore miss is recoverable in
106-
/// practice because the control plane will re-emit the agg_config
107-
/// on its next reconcile pass.
103+
/// Write one PrecomputedOutput; false reports a rejected output to the caller.
104+
/// Missing configuration or incompatible state must surface as failure:
105+
/// a finite-input completion barrier cannot acknowledge dropped outputs.
108106
///
109107
/// PR-6 follow-up: resolves the source `AggregationConfig` via
110108
/// `PolicyRegistry::get(output.policy_fp)`. The legacy
@@ -160,8 +158,18 @@ impl OutputSink for SketchStoreSink {
160158
return Ok(());
161159
}
162160
let _span = debug_span!("sketch_index_insert", batch_size = outputs.len()).entered();
161+
let mut failed = 0;
163162
for (output, accumulator) in &outputs {
164-
self.append_to_index(output, accumulator.as_ref());
163+
if !self.append_to_index(output, accumulator.as_ref()) {
164+
failed += 1;
165+
}
166+
}
167+
if failed > 0 {
168+
return Err(format!(
169+
"SketchStore rejected {failed} of {} completed outputs",
170+
outputs.len()
171+
)
172+
.into());
165173
}
166174
Ok(())
167175
}
@@ -375,9 +383,9 @@ mod tests {
375383
}
376384

377385
#[test]
378-
fn sketch_index_sink_skips_unknown_agg_id_gracefully() {
386+
fn sketch_index_sink_reports_unknown_policy_as_failure() {
379387
// Streaming config does NOT contain agg_id=99 — the sink
380-
// skips it (warn log) rather than panicking.
388+
// reports a recoverable error rather than acknowledging a lost write.
381389
let streaming = StreamingConfig::new(HashMap::new());
382390
let hot_reload = HotReloadStreamingConfig::new(streaming.clone());
383391
let sketch_index = Arc::new(SketchStore::new());
@@ -389,12 +397,13 @@ mod tests {
389397

390398
let output = PrecomputedOutput::new(1000, 2000, None, asap_types::PolicyFingerprint(99));
391399
let acc: Box<dyn AggregateCore> = Box::new(SumAccumulator::with_sum(1.0));
392-
sink.emit_batch(vec![(output, acc)]).expect("emit ok");
400+
sink.emit_batch(vec![(output, acc)])
401+
.expect_err("unpersisted output must not be acknowledged");
393402
assert_eq!(sketch_index.instance_count(), 0);
394403
}
395404

396405
/// CQ-6 — a registry-miss (policy_fp not in the running streaming
397-
/// config) is a silent write-skip; with an `IngestObservability`
406+
/// config) reports a failed write; with an `IngestObservability`
398407
/// handle wired in, the `dropped_policy_miss` counter must tick.
399408
#[test]
400409
fn sink_increments_policy_miss_counter_on_registry_miss() {
@@ -412,7 +421,8 @@ mod tests {
412421
// policy_fp=42 is absent from the empty registry → registry miss.
413422
let output = PrecomputedOutput::new(1000, 2000, None, asap_types::PolicyFingerprint(42));
414423
let acc: Box<dyn AggregateCore> = Box::new(SumAccumulator::with_sum(1.0));
415-
sink.emit_batch(vec![(output, acc)]).expect("emit ok");
424+
sink.emit_batch(vec![(output, acc)])
425+
.expect_err("unpersisted output must not be acknowledged");
416426

417427
assert_eq!(
418428
sketch_index.instance_count(),
@@ -430,7 +440,8 @@ mod tests {
430440
// miss — it must not bump the counter.
431441
let unset = PrecomputedOutput::new(1000, 2000, None, asap_types::PolicyFingerprint::UNSET);
432442
let acc2: Box<dyn AggregateCore> = Box::new(SumAccumulator::with_sum(1.0));
433-
sink.emit_batch(vec![(unset, acc2)]).expect("emit ok");
443+
sink.emit_batch(vec![(unset, acc2)])
444+
.expect_err("unpersisted output must not be acknowledged");
434445
assert_eq!(
435446
obs.dropped_policy_miss
436447
.load(std::sync::atomic::Ordering::Relaxed),

‎data_plane/src/precompute_engine/series_router.rs‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,8 @@ pub enum WorkerMessage {
8585
},
8686
/// Signal the worker to flush/check idle windows.
8787
Flush,
88+
/// Finite-input barrier: acknowledge only after queued input and trailing panes reach the sink.
89+
Drain(tokio::sync::oneshot::Sender<Result<(), String>>),
8890
/// Graceful shutdown.
8991
Shutdown,
9092
}
@@ -126,6 +128,7 @@ impl fmt::Debug for WorkerMessage {
126128
.field("accumulator_type", &accumulator.type_name())
127129
.finish(),
128130
Self::Flush => f.write_str("Flush"),
131+
Self::Drain(_) => f.write_str("Drain"),
129132
Self::Shutdown => f.write_str("Shutdown"),
130133
}
131134
}
@@ -205,7 +208,7 @@ impl SeriesRouter {
205208
WorkerMessage::GroupSamples { sid, .. }
206209
| WorkerMessage::AccumulatorInput { sid, .. } => self.worker_for_sid(*sid),
207210
WorkerMessage::RawSamples { series_key, .. } => self.worker_for(series_key),
208-
WorkerMessage::Flush | WorkerMessage::Shutdown => 0,
211+
WorkerMessage::Flush | WorkerMessage::Drain(_) | WorkerMessage::Shutdown => 0,
209212
};
210213
let permit = self.senders[worker_idx]
211214
.clone()
@@ -244,6 +247,23 @@ impl SeriesRouter {
244247
Ok(())
245248
}
246249

250+
/// Caller must stop new input before invoking this finite-source barrier.
251+
pub async fn drain(&self) -> Result<(), String> {
252+
let mut replies = Vec::new();
253+
for sender in &self.senders {
254+
let (tx, rx) = tokio::sync::oneshot::channel();
255+
sender
256+
.send(WorkerMessage::Drain(tx))
257+
.await
258+
.map_err(|e| e.to_string())?;
259+
replies.push(rx);
260+
}
261+
for reply in replies {
262+
reply.await.map_err(|e| e.to_string())??;
263+
}
264+
Ok(())
265+
}
266+
247267
/// Determine which worker handles a given sid bucket.
248268
///
249269
/// Hashes the sid alone — the legacy `(agg_id, group_key)` tuple folded

0 commit comments

Comments
 (0)