aboutsummaryrefslogtreecommitdiff
path: root/src/main/kotlin/moe/nea/ledger/utils/telemetry/Span.kt
blob: 0d680a9835b66d42fc2de71c85601dc8f770adb4 (plain)
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
package moe.nea.ledger.utils.telemetry

class Span(val parent: Span?) : AutoCloseable {
	companion object {
		private val _current = object : InheritableThreadLocal<Span>() {
			override fun initialValue(): Span {
				return Span(null)
			}

			override fun childValue(parentValue: Span?): Span {
				return parentValue?.forkNewRoot() ?: initialValue()
			}
		}

		fun current(): Span {
			return _current.get()
		}
	}

	private val data = Context()

	// TODO : replace string key with a SpanKey<T> class
	fun add(key: String, value: ContextValue) {
		data.add(key, value)
	}

	/**
	 * Create a sub span, and [enter] it, with the given values.
	 */
	fun <T> enterWith(vararg pairs: Pair<String, ContextValue>, block: Span.() -> T): T {
		return enter().use { span ->
			pairs.forEach { (k, value) ->
				span.add(k, value)
			}
			block(span)
		}
	}

	/**
	 * Create a sub span, to attach some additional context, without modifying the [current] at all.
	 */
	fun forkWith(vararg pairs: Pair<String, ContextValue?>): Span {
		val newSpan = fork()
		for ((key, value) in pairs) {
			if (value == null) continue
			newSpan.add(key, value)
		}
		return newSpan
	}

	/**
	 * Create a sub span, to which additional context can be added. This context will receive updates from its parent,
	 * and will be set as the [current]. To return to the parent, either call [exit] on the child. Or use inside of a
	 * [use] block.
	 */
	fun enter(): Span {
		require(_current.get() == this)
		return fork().enterSelf()
	}

	/**
	 * Force [enter] this span, without creating a subspan. This bypasses checks like parent / child being the [current].
	 */
	fun enterSelf(): Span {
		_current.set(this)
		return this
	}

	/**
	 * Creates a temporary sub span, to which additional context can be added. This context will receive updates from
	 * its parent, but will not be set as the [current].
	 */
	fun fork(): Span {
		return Span(this)
	}

	/**
	 * Create a new root span, that will not receive any updates from the current span, but will have all the same
	 * context keys associated.
	 */
	fun forkNewRoot(): Span {
		val newRoot = Span(null)
		newRoot.data.data.putAll(collectContext().data)
		return newRoot
	}

	/**
	 * Collect the context, including all parent context
	 */
	fun collectContext(): Context {
		if (parent != null)
			return data.combineWith(parent.collectContext())
		return data
	}

	/**
	 * Exit an [entered][enter] span, returning back to the parent context, and discard any current keys.
	 */
	fun exit() {
		require(parent != null)
		require(_current.get() == this)
		_current.set(parent)
	}

	/**
	 * [AutoCloseable] implementation for [exit]
	 */
	override fun close() {
		return exit()
	}

	/**
	 * Record an empty event given the context. This indicates nothing except for "I was here".
	 * @see recordMessageEvent
	 * @see recordException
	 */
	fun recordEmptyTrace(recorder: EventRecorder) {
		recorder.record(RecordedEvent(collectContext()))
	}

	/**
	 * Record a message with the key `"event_message"` to the recorder
	 */
	fun recordMessageEvent(
		recorder: EventRecorder,
		message: String
	) {
		forkWith(CommonKeys.EVENT_MESSAGE to ContextValue.string(message))
			.recordEmptyTrace(recorder)
	}

	/**
	 * Record an exception to the recorder
	 */
	fun recordException(
		recorder: EventRecorder,
		exception: Throwable,
		message: String? = null
	) {
		forkWith(
			CommonKeys.EVENT_MESSAGE to message?.let(ContextValue::string),
			CommonKeys.EXCEPTION to ExceptionContextValue(exception),
		).recordEmptyTrace(recorder)
	}

}