|
| 1 | +// Licensed to the Apache Software Foundation (ASF) under one or more |
| 2 | +// contributor license agreements. See the NOTICE file distributed with |
| 3 | +// this work for additional information regarding copyright ownership. |
| 4 | +// The ASF licenses this file to You under the Apache License, Version 2.0 |
| 5 | +// (the "License"); you may not use this file except in compliance with |
| 6 | +// the License. You may obtain a copy of the License at |
| 7 | +// |
| 8 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +// |
| 10 | +// Unless required by applicable law or agreed to in writing, software |
| 11 | +// distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | +// See the License for the specific language governing permissions and |
| 14 | +// limitations under the License. |
| 15 | + |
| 16 | +use super::{Plugin, log_exception}; |
| 17 | +use crate::{ |
| 18 | + component::COMPONENT_AMQP_PRODUCER_ID, |
| 19 | + context::{RequestContext, SW_HEADER}, |
| 20 | + execute::{AfterExecuteHook, BeforeExecuteHook, get_this_mut, validate_num_args}, |
| 21 | + tag::{TAG_MQ_BROKER, TAG_MQ_QUEUE, TAG_MQ_TOPIC}, |
| 22 | +}; |
| 23 | +use phper::{ |
| 24 | + arrays::ZArray, |
| 25 | + objects::ZObj, |
| 26 | + values::{ExecuteData, ZVal}, |
| 27 | +}; |
| 28 | +use skywalking::{ |
| 29 | + proto::v3::SpanLayer, |
| 30 | + trace::span::{HandleSpanObject, Span}, |
| 31 | +}; |
| 32 | + |
| 33 | +#[derive(Default, Clone)] |
| 34 | +pub struct AmqpPlugin; |
| 35 | + |
| 36 | +impl Plugin for AmqpPlugin { |
| 37 | + fn class_names(&self) -> Option<&'static [&'static str]> { |
| 38 | + Some(&["AMQPExchange"]) |
| 39 | + } |
| 40 | + |
| 41 | + fn function_name_prefix(&self) -> Option<&'static str> { |
| 42 | + None |
| 43 | + } |
| 44 | + |
| 45 | + fn hook( |
| 46 | + &self, class_name: Option<&str>, function_name: &str, |
| 47 | + ) -> Option<( |
| 48 | + Box<crate::execute::BeforeExecuteHook>, |
| 49 | + Box<crate::execute::AfterExecuteHook>, |
| 50 | + )> { |
| 51 | + match (class_name, function_name) { |
| 52 | + (Some(class_name @ "AMQPExchange"), function_name @ "publish") => { |
| 53 | + Some(self.hook_exchange_publish(class_name, function_name)) |
| 54 | + } |
| 55 | + _ => None, |
| 56 | + } |
| 57 | + } |
| 58 | +} |
| 59 | + |
| 60 | +impl AmqpPlugin { |
| 61 | + fn hook_exchange_publish( |
| 62 | + &self, class_name: &str, function_name: &str, |
| 63 | + ) -> (Box<BeforeExecuteHook>, Box<AfterExecuteHook>) { |
| 64 | + let class_name = class_name.to_owned(); |
| 65 | + let function_name = function_name.to_owned(); |
| 66 | + ( |
| 67 | + Box::new(move |request_id, execute_data| { |
| 68 | + validate_num_args(execute_data, 1)?; |
| 69 | + |
| 70 | + let this = get_this_mut(execute_data)?; |
| 71 | + |
| 72 | + let peer = Self::get_peer(this)?; |
| 73 | + |
| 74 | + let exchange = this |
| 75 | + .call("getName", []) |
| 76 | + .ok() |
| 77 | + .and_then(|v| { |
| 78 | + v.as_z_str() |
| 79 | + .and_then(|s| s.to_str().ok()) |
| 80 | + .map(ToOwned::to_owned) |
| 81 | + }) |
| 82 | + .unwrap_or_default(); |
| 83 | + |
| 84 | + let routing_key = if execute_data.num_args() >= 2 { |
| 85 | + execute_data |
| 86 | + .get_parameter(1) |
| 87 | + .as_z_str() |
| 88 | + .and_then(|s| s.to_str().ok()) |
| 89 | + .map(ToOwned::to_owned) |
| 90 | + .unwrap_or_default() |
| 91 | + } else { |
| 92 | + String::new() |
| 93 | + }; |
| 94 | + |
| 95 | + let span = Self::create_exit_span( |
| 96 | + request_id, |
| 97 | + &class_name, |
| 98 | + &function_name, |
| 99 | + &peer, |
| 100 | + &exchange, |
| 101 | + &routing_key, |
| 102 | + )?; |
| 103 | + |
| 104 | + Self::inject_sw_header(request_id, execute_data, &peer)?; |
| 105 | + |
| 106 | + Ok(Box::new(span)) |
| 107 | + }), |
| 108 | + Box::new(move |_, span, _, _| { |
| 109 | + let mut span = span.downcast::<Span>().unwrap(); |
| 110 | + log_exception(&mut *span); |
| 111 | + Ok(()) |
| 112 | + }), |
| 113 | + ) |
| 114 | + } |
| 115 | + |
| 116 | + fn get_peer(this: &mut ZObj) -> crate::Result<String> { |
| 117 | + let mut channel = this.call("getChannel", [])?; |
| 118 | + let channel = channel |
| 119 | + .expect_mut_z_obj() |
| 120 | + .map_err(|e| anyhow::anyhow!("channel isn't object: {}", e))?; |
| 121 | + let mut connection = channel.call("getConnection", [])?; |
| 122 | + let connection = connection |
| 123 | + .expect_mut_z_obj() |
| 124 | + .map_err(|e| anyhow::anyhow!("connection isn't object: {}", e))?; |
| 125 | + let host = connection.call("getHost", [])?; |
| 126 | + let host = host |
| 127 | + .expect_z_str() |
| 128 | + .map_err(|e| anyhow::anyhow!("host isn't string: {}", e))? |
| 129 | + .to_str()?; |
| 130 | + let port = connection.call("getPort", [])?; |
| 131 | + let port = port.as_long().unwrap_or_default(); |
| 132 | + Ok(format!("{}:{}", host, port)) |
| 133 | + } |
| 134 | + |
| 135 | + fn create_exit_span( |
| 136 | + request_id: Option<i64>, class_name: &str, function_name: &str, peer: &str, exchange: &str, |
| 137 | + routing_key: &str, |
| 138 | + ) -> crate::Result<Span> { |
| 139 | + let mut span = RequestContext::try_with_global_ctx(request_id, |ctx| { |
| 140 | + Ok(ctx.create_exit_span(&format!("{}->{}", class_name, function_name), peer)) |
| 141 | + })?; |
| 142 | + |
| 143 | + let span_object = span.span_object_mut(); |
| 144 | + span_object.set_span_layer(SpanLayer::Mq); |
| 145 | + span_object.component_id = COMPONENT_AMQP_PRODUCER_ID; |
| 146 | + span_object.add_tag(TAG_MQ_BROKER, peer); |
| 147 | + span_object.add_tag(TAG_MQ_TOPIC, exchange); |
| 148 | + span_object.add_tag(TAG_MQ_QUEUE, routing_key); |
| 149 | + |
| 150 | + Ok(span) |
| 151 | + } |
| 152 | + |
| 153 | + fn inject_sw_header( |
| 154 | + request_id: Option<i64>, execute_data: &mut ExecuteData, peer: &str, |
| 155 | + ) -> crate::Result<()> { |
| 156 | + let sw_header = RequestContext::try_get_sw_header(request_id, peer)?; |
| 157 | + |
| 158 | + let attributes = Self::ensure_attributes(execute_data)?; |
| 159 | + let headers = Self::ensure_headers(attributes)?; |
| 160 | + headers.insert(SW_HEADER, sw_header); |
| 161 | + |
| 162 | + Ok(()) |
| 163 | + } |
| 164 | + |
| 165 | + fn ensure_attributes( |
| 166 | + execute_data: &mut ExecuteData, |
| 167 | + ) -> crate::Result<&mut phper::arrays::ZArr> { |
| 168 | + // php-amqp parses publish() as `s|s!l!a/`, so the 4th `headers` |
| 169 | + // argument is only visible after the missing optional parameters are |
| 170 | + // materialized and `num_args` reflects that. |
| 171 | + execute_data.materialize_missing([ |
| 172 | + ZVal::from(()), |
| 173 | + ZVal::from(()), |
| 174 | + ZVal::from(ZArray::new()), |
| 175 | + ])?; |
| 176 | + |
| 177 | + let attributes = execute_data.get_mut_parameter(3); |
| 178 | + if attributes.as_z_arr().is_none() { |
| 179 | + *attributes = ZArray::new().into(); |
| 180 | + } |
| 181 | + |
| 182 | + Ok(attributes |
| 183 | + .as_mut_z_arr() |
| 184 | + .ok_or_else(|| anyhow::anyhow!("attributes isn't array"))?) |
| 185 | + } |
| 186 | + |
| 187 | + fn ensure_headers( |
| 188 | + attributes: &mut phper::arrays::ZArr, |
| 189 | + ) -> crate::Result<&mut phper::arrays::ZArr> { |
| 190 | + let has_headers = attributes |
| 191 | + .get("headers") |
| 192 | + .and_then(|headers| headers.as_z_arr()) |
| 193 | + .is_some(); |
| 194 | + if !has_headers { |
| 195 | + attributes.insert("headers", ZArray::new()); |
| 196 | + } |
| 197 | + |
| 198 | + Ok(attributes |
| 199 | + .get_mut("headers") |
| 200 | + .and_then(|headers| headers.as_mut_z_arr()) |
| 201 | + .ok_or_else(|| anyhow::anyhow!("headers isn't array"))?) |
| 202 | + } |
| 203 | +} |
0 commit comments