1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165

// Copyright 2020 Parity Technologies (UK) Ltd.
// This file is part of Substrate.

// Substrate is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Substrate is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Substrate.  If not, see <http://www.gnu.org/licenses/>.

//! Proxy to allow entering tracing spans from wasm.
//!
//! Use `enter_span` and `exit_span` to surround the code that you wish to trace
use rental;
use tracing::info_span;

/// Used to identify a proxied WASM trace
pub const WASM_TRACE_IDENTIFIER: &'static str = "WASM_TRACE";
/// Used to extract the real `target` from the associated values of the span
pub const WASM_TARGET_KEY: &'static str = "proxied_wasm_target";
/// Used to extract the real `name` from the associated values of the span
pub const WASM_NAME_KEY: &'static str = "proxied_wasm_name";

const MAX_SPANS_LEN: usize = 1000;

rental! {
	pub mod rent_span {
		#[rental]
		pub struct SpanAndGuard {
			span: Box<tracing::Span>,
			guard: tracing::span::Entered<'span>,
		}
	}
}

/// Requires a tracing::Subscriber to process span traces,
/// this is available when running with client (and relevant cli params).
pub struct TracingProxy {
	next_id: u64,
	spans: Vec<(u64, rent_span::SpanAndGuard)>,
}

impl Drop for TracingProxy {
	fn drop(&mut self) {
		if !self.spans.is_empty() {
			log::debug!(
				target: "tracing",
				"Dropping TracingProxy with {} un-exited spans, marking as not valid", self.spans.len()
			);
			while let Some((_, mut sg)) = self.spans.pop() {
				sg.rent_all_mut(|s| { s.span.record("is_valid_trace", &false); });
			}
		}
	}
}

impl TracingProxy {
	pub fn new() -> TracingProxy {
		TracingProxy {
			next_id: 0,
			spans: Vec::new(),
		}
	}
}

impl TracingProxy {
	/// Create and enter a `tracing` Span, returning the span id,
	/// which should be passed to `exit_span(id)` to signal that the span should exit.
	pub fn enter_span(&mut self, proxied_wasm_target: &str, proxied_wasm_name: &str) -> u64 {
		// The identifiers `proxied_wasm_target` and `proxied_wasm_name` must match their associated const,
		// WASM_TARGET_KEY and WASM_NAME_KEY.
		let span = info_span!(WASM_TRACE_IDENTIFIER, is_valid_trace = true, proxied_wasm_target, proxied_wasm_name);
		self.next_id += 1;
		let sg = rent_span::SpanAndGuard::new(
			Box::new(span),
			|span| span.enter(),
		);
		self.spans.push((self.next_id, sg));
		if self.spans.len() > MAX_SPANS_LEN {
			// This is to prevent unbounded growth of Vec and could mean one of the following:
			// 1. Too many nested spans, or MAX_SPANS_LEN is too low.
			// 2. Not correctly exiting spans due to misconfiguration / misuse
			log::warn!(
				target: "tracing",
				"TracingProxy MAX_SPANS_LEN exceeded, removing oldest span."
			);
			let mut sg = self.spans.remove(0).1;
			sg.rent_all_mut(|s| { s.span.record("is_valid_trace", &false); });
		}
		self.next_id
	}

	/// Exit a span by dropping it along with it's associated guard.
	pub fn exit_span(&mut self, id: u64) {
		if self.spans.last().map(|l| id > l.0).unwrap_or(true) {
			log::warn!(target: "tracing", "Span id not found in TracingProxy: {}", id);
			return;
		}
		let mut last_span = self.spans.pop().expect("Just checked that there is an element to pop; qed");
		while id < last_span.0 {
			log::warn!(
				target: "tracing",
				"TracingProxy Span ids not equal! id parameter given: {}, last span: {}",
				id,
				last_span.0,
			);
			last_span.1.rent_all_mut(|s| { s.span.record("is_valid_trace", &false); });
			if let Some(s) = self.spans.pop() {
				last_span = s;
			} else {
				log::warn!(target: "tracing", "Span id not found in TracingProxy {}", id);
				return;
			}
		}
	}
}


#[cfg(test)]
mod tests {
	use super::*;

	fn create_spans(proxy: &mut TracingProxy, qty: usize) -> Vec<u64> {
		let mut spans = Vec::new();
		for n in 0..qty {
			spans.push(proxy.enter_span("target", &format!("{}", n)));
		}
		spans
	}

	#[test]
	fn max_spans_len_respected() {
		let mut proxy = TracingProxy::new();
		let _spans = create_spans(&mut proxy, MAX_SPANS_LEN + 10);
		assert_eq!(proxy.spans.len(), MAX_SPANS_LEN);
		// ensure oldest spans removed
		assert_eq!(proxy.spans[0].0, 11);
	}

	#[test]
	fn handles_span_exit_scenarios() {
		let mut proxy = TracingProxy::new();
		let _spans = create_spans(&mut proxy, 10);
		assert_eq!(proxy.spans.len(), 10);
		// exit span normally
		proxy.exit_span(10);
		assert_eq!(proxy.spans.len(), 9);
		// skip and exit outer span without exiting inner, id: 8 instead of 9
		proxy.exit_span(8);
		// should have also removed the inner span that was lost
		assert_eq!(proxy.spans.len(), 7);
		// try to exit span not held
		proxy.exit_span(9);
		assert_eq!(proxy.spans.len(), 7);
		// exit all spans
		proxy.exit_span(1);
		assert_eq!(proxy.spans.len(), 0);
	}
}