1use crate::llm::actions::protocol_trait::{ActionResult, Protocol, Server};
14use crate::llm::actions::{ActionDefinition, Parameter};
15use crate::protocol::log_template::LogTemplate;
16use crate::protocol::EventType;
17use crate::state::app_state::AppState;
18use anyhow::Result;
19use serde_json::{json, Value};
20use std::sync::LazyLock;
21
22pub struct YarnProtocol {}
24
25impl YarnProtocol {
26 pub fn new() -> Self {
27 Self {}
28 }
29}
30
31impl Default for YarnProtocol {
32 fn default() -> Self {
33 Self::new()
34 }
35}
36
37pub static YARN_REQUEST_EVENT: LazyLock<EventType> = LazyLock::new(|| {
40 EventType::new(
41 "yarn_request",
42 "Hadoop YARN ResourceManager REST request received",
43 json!({
44 "type": "send_yarn_apps",
45 "apps": [{
46 "id": "application_1476912658570_0002",
47 "user": "dr.who",
48 "name": "word count",
49 "queue": "default",
50 "state": "FINISHED",
51 "finalStatus": "SUCCEEDED",
52 "progress": 100.0,
53 "applicationType": "MAPREDUCE"
54 }]
55 }),
56 )
57 .with_parameters(vec![
58 Parameter {
59 name: "method".to_string(),
60 type_hint: "string".to_string(),
61 description: "HTTP method (GET, POST)".to_string(),
62 required: true,
63 },
64 Parameter {
65 name: "path".to_string(),
66 type_hint: "string".to_string(),
67 description: "Request path, e.g. /ws/v1/cluster/apps".to_string(),
68 required: true,
69 },
70 Parameter {
71 name: "operation".to_string(),
72 type_hint: "string".to_string(),
73 description: "Detected operation: metrics, scheduler, apps, submit, \
74 new_application, app, nodes, node"
75 .to_string(),
76 required: true,
77 },
78 Parameter {
79 name: "app_id".to_string(),
80 type_hint: "string".to_string(),
81 description: "Application id from the path (for the 'app' operation)".to_string(),
82 required: false,
83 },
84 Parameter {
85 name: "node_id".to_string(),
86 type_hint: "string".to_string(),
87 description: "Node id from the path (for the 'node' operation)".to_string(),
88 required: false,
89 },
90 Parameter {
91 name: "request_body".to_string(),
92 type_hint: "string".to_string(),
93 description: "JSON request body (for application submission)".to_string(),
94 required: false,
95 },
96 ])
97 .with_actions(yarn_actions())
98 .with_log_template(
99 LogTemplate::new()
100 .with_info("YARN {client_ip} {method} {path}")
101 .with_debug("YARN {method} {path} op={operation} from {client_ip}:{client_port}")
102 .with_trace("YARN: {json_pretty(.)}"),
103 )
104});
105
106fn yarn_actions() -> Vec<ActionDefinition> {
108 vec![
109 send_yarn_metrics_action(),
110 send_yarn_scheduler_action(),
111 send_yarn_apps_action(),
112 send_yarn_app_action(),
113 send_yarn_nodes_action(),
114 send_yarn_new_application_action(),
115 send_yarn_submit_response_action(),
116 send_yarn_error_action(),
117 ]
118}
119
120fn send_yarn_metrics_action() -> ActionDefinition {
121 ActionDefinition {
122 name: "send_yarn_metrics".to_string(),
123 description: "Answer GET /ws/v1/cluster/metrics with the cluster's aggregate metrics. \
124 All fields are numbers describing the (invented) cluster; keep them \
125 self-consistent (allocatedMB <= totalMB, activeNodes <= totalNodes). \
126 Wrapped in the {\"clusterMetrics\": {...}} envelope automatically."
127 .to_string(),
128 parameters: vec![Parameter {
129 name: "metrics".to_string(),
130 type_hint: "object".to_string(),
131 description: "Object of clusterMetrics fields: appsSubmitted, appsRunning, \
132 appsPending, appsCompleted, appsFailed, appsKilled, totalMB, \
133 availableMB, allocatedMB, reservedMB, totalVirtualCores, \
134 availableVirtualCores, allocatedVirtualCores, containersAllocated, \
135 totalNodes, activeNodes, lostNodes, unhealthyNodes, \
136 decommissionedNodes. Any field you omit defaults to 0."
137 .to_string(),
138 required: true,
139 }],
140 example: json!({
141 "type": "send_yarn_metrics",
142 "metrics": {
143 "appsSubmitted": 4, "appsRunning": 1, "appsCompleted": 3,
144 "totalMB": 32768, "availableMB": 24576, "allocatedMB": 8192,
145 "totalVirtualCores": 16, "availableVirtualCores": 12, "allocatedVirtualCores": 4,
146 "containersAllocated": 4, "totalNodes": 2, "activeNodes": 2
147 }
148 }),
149 log_template: Some(
150 LogTemplate::new().with_info("-> YARN metrics {metrics.totalNodes} nodes"),
151 ),
152 }
153}
154
155fn send_yarn_scheduler_action() -> ActionDefinition {
156 ActionDefinition {
157 name: "send_yarn_scheduler".to_string(),
158 description: "Answer GET /ws/v1/cluster/scheduler with the scheduler and its queue tree. \
159 Rendered in the CapacityScheduler's own shape, {\"scheduler\": \
160 {\"schedulerInfo\": {\"type\": \"capacityScheduler\", \"queueName\": \"root\", \
161 ..., \"queues\": {\"queue\": [...]}}}}, with each queue's children nested the \
162 same way; or, with scheduler_type fifoScheduler, the FifoScheduler's flat \
163 summary. Keep it consistent with the apps you report (a queue's \
164 numApplications)."
165 .to_string(),
166 parameters: vec![
167 Parameter {
168 name: "scheduler_type".to_string(),
169 type_hint: "string".to_string(),
170 description: "capacityScheduler (default) or fifoScheduler".to_string(),
171 required: false,
172 },
173 Parameter {
174 name: "root".to_string(),
175 type_hint: "object".to_string(),
176 description: "The root queue's (or the FIFO scheduler's) own numbers: capacity \
177 (default 100), usedCapacity (default 0), maxCapacity (default \
178 100); for fifoScheduler also qstate, numNodes, numContainers, \
179 usedNodeCapacity, availNodeCapacity, totalNodeCapacity, \
180 minQueueMemoryCapacity, maxQueueMemoryCapacity."
181 .to_string(),
182 required: false,
183 },
184 Parameter {
185 name: "queues".to_string(),
186 type_hint: "array".to_string(),
187 description: format!(
188 "capacityScheduler only, required: the queues under root. Each is an \
189 object with queueName (required) and capacity, usedCapacity, maxCapacity, \
190 absoluteCapacity, absoluteUsedCapacity, absoluteMaxCapacity, state \
191 (RUNNING/STOPPED), numApplications, numActiveApplications, \
192 numPendingApplications, numContainers, maxApplications, userLimit, \
193 userLimitFactor; a parent queue carries its children as a `queues` array \
194 of the same objects, at most {MAX_QUEUE_DEPTH} levels below root. \
195 queuePath is filled in from the names."
196 ),
197 required: false,
198 },
199 ],
200 example: json!({
201 "type": "send_yarn_scheduler",
202 "root": {"capacity": 100.0, "usedCapacity": 25.0, "maxCapacity": 100.0},
203 "queues": [
204 {"queueName": "default", "capacity": 60.0, "usedCapacity": 0.0,
205 "maxCapacity": 100.0, "numApplications": 0},
206 {"queueName": "etl", "capacity": 40.0, "usedCapacity": 62.5,
207 "maxCapacity": 100.0, "numApplications": 1,
208 "queues": [{"queueName": "nightly", "capacity": 100.0,
209 "usedCapacity": 62.5, "numApplications": 1}]}
210 ]
211 }),
212 log_template: Some(LogTemplate::new().with_info("-> YARN scheduler {queues_len} queues")),
213 }
214}
215
216fn send_yarn_apps_action() -> ActionDefinition {
217 ActionDefinition {
218 name: "send_yarn_apps".to_string(),
219 description: "Answer GET /ws/v1/cluster/apps with the list of applications. Supply the \
220 `apps` array (empty array = no applications, rendered as the YARN idiom \
221 \"apps\": null). Each element is an application object; wrapped as \
222 {\"apps\": {\"app\": [...]}} automatically."
223 .to_string(),
224 parameters: vec![Parameter {
225 name: "apps".to_string(),
226 type_hint: "array".to_string(),
227 description: "Array of application objects. Common fields per app: id \
228 (e.g. application_1476912658570_0002), user, name, queue, state \
229 (NEW/SUBMITTED/ACCEPTED/RUNNING/FINISHED/FAILED/KILLED), finalStatus \
230 (UNDEFINED/SUCCEEDED/FAILED/KILLED), progress (0-100), \
231 applicationType (MAPREDUCE/SPARK/...), startedTime, finishedTime, \
232 elapsedTime, allocatedMB, allocatedVCores, runningContainers."
233 .to_string(),
234 required: true,
235 }],
236 example: json!({
237 "type": "send_yarn_apps",
238 "apps": [{
239 "id": "application_1476912658570_0002", "user": "dr.who", "name": "word count",
240 "queue": "default", "state": "RUNNING", "finalStatus": "UNDEFINED",
241 "progress": 62.5, "applicationType": "MAPREDUCE", "allocatedMB": 4096,
242 "allocatedVCores": 2, "runningContainers": 2
243 }]
244 }),
245 log_template: Some(LogTemplate::new().with_info("-> YARN {apps_len} apps")),
246 }
247}
248
249fn send_yarn_app_action() -> ActionDefinition {
250 ActionDefinition {
251 name: "send_yarn_app".to_string(),
252 description: "Answer GET /ws/v1/cluster/apps/{appid} with a single application. Supply \
253 the `app` object; wrapped as {\"app\": {...}} automatically. Emit \
254 send_yarn_error with status 404 if the application does not exist."
255 .to_string(),
256 parameters: vec![Parameter {
257 name: "app".to_string(),
258 type_hint: "object".to_string(),
259 description: "A single application object (same fields as an element of \
260 send_yarn_apps.apps)."
261 .to_string(),
262 required: true,
263 }],
264 example: json!({
265 "type": "send_yarn_app",
266 "app": {
267 "id": "application_1476912658570_0002", "user": "dr.who", "name": "word count",
268 "queue": "default", "state": "FINISHED", "finalStatus": "SUCCEEDED",
269 "progress": 100.0, "applicationType": "MAPREDUCE"
270 }
271 }),
272 log_template: Some(LogTemplate::new().with_info("-> YARN app {app.id}")),
273 }
274}
275
276fn send_yarn_nodes_action() -> ActionDefinition {
277 ActionDefinition {
278 name: "send_yarn_nodes".to_string(),
279 description: "Answer GET /ws/v1/cluster/nodes with the NodeManager list. Supply the \
280 `nodes` array; wrapped as {\"nodes\": {\"node\": [...]}} automatically."
281 .to_string(),
282 parameters: vec![Parameter {
283 name: "nodes".to_string(),
284 type_hint: "array".to_string(),
285 description: "Array of node objects. Common fields per node: id (host:port, e.g. \
286 host.domain.com:8041), nodeHostName, nodeHTTPAddress, rack \
287 (/default-rack), state (NEW/RUNNING/UNHEALTHY/DECOMMISSIONED/LOST), \
288 healthReport, numContainers, usedMemoryMB, availMemoryMB, \
289 usedVirtualCores, availableVirtualCores, version, lastHealthUpdate."
290 .to_string(),
291 required: true,
292 }],
293 example: json!({
294 "type": "send_yarn_nodes",
295 "nodes": [{
296 "id": "host1.example.com:8041", "nodeHostName": "host1.example.com",
297 "nodeHTTPAddress": "host1.example.com:8042", "rack": "/default-rack",
298 "state": "RUNNING", "healthReport": "", "numContainers": 2,
299 "usedMemoryMB": 4096, "availMemoryMB": 12288, "usedVirtualCores": 2,
300 "availableVirtualCores": 6, "version": "3.3.6"
301 }]
302 }),
303 log_template: Some(LogTemplate::new().with_info("-> YARN {nodes_len} nodes")),
304 }
305}
306
307fn send_yarn_new_application_action() -> ActionDefinition {
308 ActionDefinition {
309 name: "send_yarn_new_application".to_string(),
310 description: "Answer POST /ws/v1/cluster/apps/new-application (the first step of the \
311 submission flow) with a freshly allocated application id and the cluster's \
312 maximum resource capability."
313 .to_string(),
314 parameters: vec![
315 Parameter {
316 name: "application_id".to_string(),
317 type_hint: "string".to_string(),
318 description: "Newly allocated application id, e.g. \
319 application_1476912658570_0005."
320 .to_string(),
321 required: true,
322 },
323 Parameter {
324 name: "max_memory_mb".to_string(),
325 type_hint: "number".to_string(),
326 description: "maximum-resource-capability memory in MB (default 8192)".to_string(),
327 required: false,
328 },
329 Parameter {
330 name: "max_vcores".to_string(),
331 type_hint: "number".to_string(),
332 description: "maximum-resource-capability vCores (default 4)".to_string(),
333 required: false,
334 },
335 ],
336 example: json!({
337 "type": "send_yarn_new_application",
338 "application_id": "application_1476912658570_0005",
339 "max_memory_mb": 8192,
340 "max_vcores": 4
341 }),
342 log_template: Some(
343 LogTemplate::new().with_info("-> YARN new-application {application_id}"),
344 ),
345 }
346}
347
348fn send_yarn_submit_response_action() -> ActionDefinition {
349 ActionDefinition {
350 name: "send_yarn_submit_response".to_string(),
351 description: "Answer POST /ws/v1/cluster/apps (submit an application). On acceptance YARN \
352 replies 202 Accepted with an empty body and a Location header pointing at \
353 the new application; set accepted=true. To reject the submission set \
354 accepted=false and supply a message (rendered as a 400 RemoteException)."
355 .to_string(),
356 parameters: vec![
357 Parameter {
358 name: "accepted".to_string(),
359 type_hint: "boolean".to_string(),
360 description: "true = 202 Accepted, false = 400 rejection".to_string(),
361 required: true,
362 },
363 Parameter {
364 name: "application_id".to_string(),
365 type_hint: "string".to_string(),
366 description: "The submitted application id (used to build the Location header \
367 when accepted)."
368 .to_string(),
369 required: false,
370 },
371 Parameter {
372 name: "message".to_string(),
373 type_hint: "string".to_string(),
374 description: "Rejection reason (when accepted=false)".to_string(),
375 required: false,
376 },
377 ],
378 example: json!({
379 "type": "send_yarn_submit_response",
380 "accepted": true,
381 "application_id": "application_1476912658570_0005"
382 }),
383 log_template: Some(LogTemplate::new().with_info("-> YARN submit accepted={accepted}")),
384 }
385}
386
387fn send_yarn_error_action() -> ActionDefinition {
388 ActionDefinition {
389 name: "send_yarn_error".to_string(),
390 description: "Return a YARN RemoteException error envelope with an explicit HTTP status. \
391 Use for a missing application/node (404), a bad request (400) or any \
392 condition without a dedicated action. status must be 100-599."
393 .to_string(),
394 parameters: vec![
395 Parameter {
396 name: "status".to_string(),
397 type_hint: "number".to_string(),
398 description: "HTTP status code (400, 404, 500, ...)".to_string(),
399 required: true,
400 },
401 Parameter {
402 name: "exception".to_string(),
403 type_hint: "string".to_string(),
404 description: "Exception class name, e.g. NotFoundException, BadRequestException"
405 .to_string(),
406 required: false,
407 },
408 Parameter {
409 name: "message".to_string(),
410 type_hint: "string".to_string(),
411 description: "Human-readable error message".to_string(),
412 required: true,
413 },
414 ],
415 example: json!({
416 "type": "send_yarn_error",
417 "status": 404,
418 "exception": "NotFoundException",
419 "message": "app with id 'application_1_0009' not found"
420 }),
421 log_template: Some(LogTemplate::new().with_info("-> YARN error {status} {exception}")),
422 }
423}
424
425pub fn get_yarn_event_types() -> Vec<EventType> {
426 vec![YARN_REQUEST_EVENT.clone()]
427}
428
429pub const MAX_QUEUE_DEPTH: usize = 8;
433
434fn capacity_queue(queue: &Value, parent_path: &str, depth: usize) -> Result<Value> {
438 if depth > MAX_QUEUE_DEPTH {
439 return Err(anyhow::anyhow!(
440 "the queue tree is deeper than {MAX_QUEUE_DEPTH} levels below root"
441 ));
442 }
443 let mut out = queue
444 .as_object()
445 .cloned()
446 .ok_or_else(|| anyhow::anyhow!("each queue is an object"))?;
447 let name = out
448 .get("queueName")
449 .and_then(|v| v.as_str())
450 .filter(|n| !n.is_empty())
451 .ok_or_else(|| anyhow::anyhow!("each queue needs a queueName"))?
452 .to_string();
453 let path = format!("{parent_path}.{name}");
454 out.entry("queuePath").or_insert_with(|| json!(path));
455 out.entry("state").or_insert_with(|| json!("RUNNING"));
456 match out.remove("queues") {
457 Some(Value::Array(children)) if !children.is_empty() => {
458 let children = children
459 .iter()
460 .map(|c| capacity_queue(c, &path, depth + 1))
461 .collect::<Result<Vec<_>>>()?;
462 out.insert("queueType".into(), json!("parent"));
463 out.insert("queues".into(), json!({ "queue": children }));
464 }
465 None | Some(Value::Null) | Some(Value::Array(_)) => {
466 out.insert("queueType".into(), json!("leaf"));
467 out.entry("type")
468 .or_insert_with(|| json!("capacitySchedulerLeafQueueInfo"));
469 }
470 Some(_) => return Err(anyhow::anyhow!("a queue's `queues` is an array of queues")),
471 }
472 Ok(Value::Object(out))
473}
474
475fn body_of(value: &Value) -> String {
477 serde_json::to_string(value).unwrap_or_else(|_| value.to_string())
478}
479
480fn parse_status(action: &Value, key: &str) -> Result<u16> {
483 let raw = action
484 .get(key)
485 .and_then(|v| v.as_u64())
486 .ok_or_else(|| anyhow::anyhow!("Missing or invalid '{key}': expected a number"))?;
487 if !(100..=599).contains(&raw) {
488 return Err(anyhow::anyhow!(
489 "Invalid {key} {raw}: must be an HTTP status between 100 and 599"
490 ));
491 }
492 Ok(raw as u16)
493}
494
495fn remote_exception(exception: &str, message: &str) -> Value {
497 json!({
498 "RemoteException": {
499 "exception": exception,
500 "message": message,
501 "javaClassName": format!("org.apache.hadoop.yarn.webapp.{exception}"),
502 }
503 })
504}
505
506impl Protocol for YarnProtocol {
507 fn get_startup_parameters(&self) -> Vec<crate::llm::actions::ParameterDefinition> {
508 use crate::llm::actions::ParameterDefinition;
509 vec![
510 ParameterDefinition {
511 name: "resource_manager_version".to_string(),
512 type_hint: "string".to_string(),
513 description: "Version string reported by the static GET /ws/v1/cluster/info \
514 banner (resourceManagerVersion / hadoopVersion). Default 3.3.6."
515 .to_string(),
516 required: false,
517 example: json!("3.3.6"),
518 default: None,
519 },
520 ParameterDefinition {
521 name: "cluster_id".to_string(),
522 type_hint: "string".to_string(),
523 description: "Numeric cluster id reported in the /ws/v1/cluster/info banner \
524 (the RM start epoch-ms). Default 1476912658570."
525 .to_string(),
526 required: false,
527 example: json!("1476912658570"),
528 default: None,
529 },
530 ]
531 }
532 fn get_async_actions(&self, _state: &AppState) -> Vec<ActionDefinition> {
533 vec![]
534 }
535 fn get_sync_actions(&self) -> Vec<ActionDefinition> {
536 yarn_actions()
537 }
538 fn protocol_name(&self) -> &'static str {
539 "YARN"
540 }
541 fn get_event_types(&self) -> Vec<EventType> {
542 get_yarn_event_types()
543 }
544 fn stack_name(&self) -> &'static str {
545 "ETH>IP>TCP>HTTP>YARN"
546 }
547 fn keywords(&self) -> Vec<&'static str> {
548 vec!["yarn", "hadoop yarn", "yarn resourcemanager"]
549 }
550 fn metadata(&self) -> crate::protocol::metadata::ProtocolMetadataV2 {
551 use crate::protocol::metadata::{DevelopmentState, ProtocolMetadataV2};
552 ProtocolMetadataV2::builder()
553 .state(DevelopmentState::Experimental)
554 .well_known_port(8088)
555 .implementation("hyper v1 HTTP/1.1 server, manual YARN RM REST envelopes")
556 .llm_control("Applications, nodes and cluster metrics invented by the LLM")
557 .e2e_testing(
558 "tests/server/yarn/datadog_check_test.rs runs Datadog's YARN integration \
559 (datadog_checks.yarn, integrations-core tag yarn-8.4.2) once against the server \
560 and asserts the cluster, application, node and queue metrics, tags, service \
561 checks and version it submitted from the model's answers to /metrics, /apps, \
562 /nodes and /scheduler. e2e_test.rs checks the documented envelopes with \
563 reqwest. NetGet's own YARN client is tested against real Hadoop, not against \
564 this server",
565 )
566 .notes(
567 "Virtual cluster, no storage. GET /ws/v1/cluster/info is answered \
568 statically; metrics/scheduler/apps/nodes are LLM-driven, the scheduler in \
569 the CapacityScheduler's nested queues.queue shape (or the FifoScheduler's). \
570 Fail-closed: LLM failure returns 503/500 RemoteException, never an \
571 empty-but-200 cluster. Experimental: one independent consumer (Datadog's \
572 check), which reads the four list endpoints only; nothing third-party has \
573 driven an application or node by id, new-application or submission.",
574 )
575 .max_inbound_bytes(crate::server::yarn::MAX_REQUEST_BODY_BYTES)
576 .request_only("The YARN REST API is HTTP request/response; a server cannot send a peer anything unprompted")
577 .build()
578 }
579 fn description(&self) -> &'static str {
580 "Hadoop YARN ResourceManager REST API"
581 }
582 fn example_prompt(&self) -> &'static str {
583 "Start a Hadoop YARN ResourceManager on port 8088 with a couple of running Spark apps"
584 }
585 fn group_name(&self) -> &'static str {
586 "Application"
587 }
588
589 fn get_startup_examples(&self) -> crate::llm::actions::StartupExamples {
590 use crate::llm::actions::StartupExamples;
591
592 let script = r#"import json, sys
595data = json.load(sys.stdin)
596event = data["event"]
597if data["event_type_id"] == "yarn_request":
598 actions = [{"type": "send_yarn_metrics",
599 "metrics": {"appsRunning": 0, "availableMB": 8192}}]
600else:
601 actions = []
602print(json.dumps({"actions": actions}))"#;
603
604 StartupExamples::new(
605 json!({
606 "type": "open_server",
607 "port": 8088,
608 "base_stack": "yarn",
609 "instruction": "Hadoop YARN ResourceManager reporting a small MapReduce cluster"
610 }),
611 json!({
612 "type": "open_server",
613 "port": 8088,
614 "base_stack": "yarn",
615 "event_handlers": [{
616 "event_pattern": "yarn_request",
617 "handler": { "type": "script", "language": "python", "code": script }
618 }]
619 }),
620 json!({
621 "type": "open_server",
622 "port": 8088,
623 "base_stack": "yarn",
624 "event_handlers": [{
625 "event_pattern": "yarn_request",
626 "handler": {
627 "type": "static",
628 "actions": [{
629 "type": "send_yarn_metrics",
630 "metrics": { "totalNodes": 1, "activeNodes": 1, "totalMB": 8192,
631 "availableMB": 8192 }
632 }]
633 }
634 }]
635 }),
636 )
637 }
638}
639
640impl Server for YarnProtocol {
641 fn spawn(
642 &self,
643 ctx: crate::protocol::SpawnContext,
644 ) -> std::pin::Pin<
645 Box<dyn std::future::Future<Output = anyhow::Result<std::net::SocketAddr>> + Send>,
646 > {
647 Box::pin(async move {
648 let rm_version = ctx
650 .startup_params
651 .as_ref()
652 .and_then(|p| {
653 p.get_optional_string("resource_manager_version")
654 .ok()
655 .flatten()
656 })
657 .unwrap_or_else(|| "3.3.6".to_string());
658 let cluster_id = ctx
659 .startup_params
660 .as_ref()
661 .and_then(|p| p.get_optional_string("cluster_id").ok().flatten())
662 .unwrap_or_else(|| "1476912658570".to_string());
663
664 crate::server::yarn::YarnServer::spawn_with_llm_actions(
665 ctx.legacy_listen_addr(),
666 ctx.llm_client,
667 ctx.state,
668 ctx.status_tx,
669 ctx.server_id,
670 rm_version,
671 cluster_id,
672 )
673 .await
674 })
675 }
676
677 fn execute_action(&self, action: Value) -> Result<ActionResult> {
678 let action_type = action
679 .get("type")
680 .and_then(|v| v.as_str())
681 .ok_or_else(|| anyhow::anyhow!("Missing action type"))?;
682
683 let custom = |status: u16, body: Value| {
684 Ok(ActionResult::Custom {
685 name: "yarn_response".to_string(),
686 data: json!({ "status": status, "body": body_of(&body), "location": Value::Null }),
687 })
688 };
689
690 match action_type {
691 "send_yarn_metrics" => {
692 let metrics = action
693 .get("metrics")
694 .cloned()
695 .filter(|v| v.is_object())
696 .ok_or_else(|| anyhow::anyhow!("Missing 'metrics' object"))?;
697 let mut m = metrics;
699 for key in [
700 "appsSubmitted",
701 "appsCompleted",
702 "appsPending",
703 "appsRunning",
704 "appsFailed",
705 "appsKilled",
706 "reservedMB",
707 "availableMB",
708 "allocatedMB",
709 "totalMB",
710 "reservedVirtualCores",
711 "availableVirtualCores",
712 "allocatedVirtualCores",
713 "totalVirtualCores",
714 "containersAllocated",
715 "containersReserved",
716 "containersPending",
717 "totalNodes",
718 "activeNodes",
719 "lostNodes",
720 "unhealthyNodes",
721 "decommissionedNodes",
722 "decommissioningNodes",
723 "rebootedNodes",
724 "shutdownNodes",
725 ] {
726 if m.get(key).is_none() {
727 m[key] = json!(0);
728 }
729 }
730 custom(200, json!({ "clusterMetrics": m }))
731 }
732 "send_yarn_scheduler" => {
733 let kind = action
734 .get("scheduler_type")
735 .and_then(|v| v.as_str())
736 .unwrap_or("capacityScheduler");
737 let mut info = action
738 .get("root")
739 .and_then(|v| v.as_object())
740 .cloned()
741 .unwrap_or_default();
742 info.insert("type".into(), json!(kind));
743 for (key, default) in [
744 ("capacity", 100.0),
745 ("usedCapacity", 0.0),
746 ("maxCapacity", 100.0),
747 ] {
748 info.entry(key).or_insert_with(|| json!(default));
749 }
750 match kind {
751 "capacityScheduler" => {
752 let queues = action
753 .get("queues")
754 .and_then(|v| v.as_array())
755 .filter(|q| !q.is_empty())
756 .ok_or_else(|| {
757 anyhow::anyhow!(
758 "a capacityScheduler needs a non-empty 'queues' array"
759 )
760 })?
761 .iter()
762 .map(|q| capacity_queue(q, "root", 1))
763 .collect::<Result<Vec<_>>>()?;
764 info.insert("queueName".into(), json!("root"));
765 info.insert("queuePath".into(), json!("root"));
766 info.insert("queueType".into(), json!("parent"));
767 info.insert("queues".into(), json!({ "queue": queues }));
768 }
769 "fifoScheduler" => {
770 info.entry("qstate").or_insert_with(|| json!("RUNNING"));
771 }
772 other => {
773 return Err(anyhow::anyhow!(
774 "scheduler_type {other:?}: capacityScheduler or fifoScheduler"
775 ))
776 }
777 }
778 custom(200, json!({ "scheduler": { "schedulerInfo": info } }))
779 }
780 "send_yarn_apps" => {
781 let apps = action
782 .get("apps")
783 .and_then(|v| v.as_array())
784 .ok_or_else(|| anyhow::anyhow!("Missing 'apps' array"))?;
785 let body = if apps.is_empty() {
787 json!({ "apps": Value::Null })
788 } else {
789 json!({ "apps": { "app": apps } })
790 };
791 custom(200, body)
792 }
793 "send_yarn_app" => {
794 let app = action
795 .get("app")
796 .cloned()
797 .filter(|v| v.is_object())
798 .ok_or_else(|| anyhow::anyhow!("Missing 'app' object"))?;
799 custom(200, json!({ "app": app }))
800 }
801 "send_yarn_nodes" => {
802 let nodes = action
803 .get("nodes")
804 .and_then(|v| v.as_array())
805 .ok_or_else(|| anyhow::anyhow!("Missing 'nodes' array"))?;
806 let body = if nodes.is_empty() {
807 json!({ "nodes": Value::Null })
808 } else {
809 json!({ "nodes": { "node": nodes } })
810 };
811 custom(200, body)
812 }
813 "send_yarn_new_application" => {
814 let app_id = action
815 .get("application_id")
816 .and_then(|v| v.as_str())
817 .ok_or_else(|| anyhow::anyhow!("Missing 'application_id'"))?;
818 let mem = action
819 .get("max_memory_mb")
820 .and_then(|v| v.as_u64())
821 .unwrap_or(8192);
822 let vcores = action
823 .get("max_vcores")
824 .and_then(|v| v.as_u64())
825 .unwrap_or(4);
826 custom(
827 200,
828 json!({
829 "application-id": app_id,
830 "maximum-resource-capability": { "memory": mem, "vCores": vcores }
831 }),
832 )
833 }
834 "send_yarn_submit_response" => {
835 let accepted = action
836 .get("accepted")
837 .and_then(|v| v.as_bool())
838 .ok_or_else(|| anyhow::anyhow!("Missing 'accepted' boolean"))?;
839 if accepted {
840 let app_id = action
841 .get("application_id")
842 .and_then(|v| v.as_str())
843 .unwrap_or("");
844 Ok(ActionResult::Custom {
846 name: "yarn_response".to_string(),
847 data: json!({
848 "status": 202,
849 "body": "",
850 "location": format!("/ws/v1/cluster/apps/{app_id}"),
851 }),
852 })
853 } else {
854 let message = action
855 .get("message")
856 .and_then(|v| v.as_str())
857 .unwrap_or("application submission rejected");
858 custom(400, remote_exception("BadRequestException", message))
859 }
860 }
861 "send_yarn_error" => {
862 let status = parse_status(&action, "status")?;
863 let message = action
864 .get("message")
865 .and_then(|v| v.as_str())
866 .ok_or_else(|| anyhow::anyhow!("Missing 'message'"))?;
867 let exception = action
868 .get("exception")
869 .and_then(|v| v.as_str())
870 .unwrap_or("WebApplicationException");
871 custom(status, remote_exception(exception, message))
872 }
873 _ => Err(anyhow::anyhow!("Unknown action type: {}", action_type)),
874 }
875 }
876}