|
| 1 | +/* |
| 2 | + * Copyright 2026, Datadog, Inc |
| 3 | + * |
| 4 | + * Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | + * you may not use this file except in compliance with the License. |
| 6 | + * 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 | + |
| 17 | +#include "context_api.h" |
| 18 | +#include "context.h" |
| 19 | +#include "guards.h" |
| 20 | +#include "otel_context.h" |
| 21 | +#include "otel_process_ctx.h" |
| 22 | +#include "profiler.h" |
| 23 | +#include "thread.h" |
| 24 | +#include <cstring> |
| 25 | + |
| 26 | +// Static member initialization |
| 27 | +ContextStorageMode ContextApi::_mode = CTX_STORAGE_PROFILER; |
| 28 | +bool ContextApi::_initialized = false; |
| 29 | +char* ContextApi::_attribute_keys[MAX_ATTRIBUTE_KEYS] = {}; |
| 30 | +int ContextApi::_attribute_key_count = 0; |
| 31 | + |
| 32 | +bool ContextApi::initialize(const Arguments& args) { |
| 33 | + if (__atomic_load_n(&_initialized, __ATOMIC_ACQUIRE)) { |
| 34 | + return true; |
| 35 | + } |
| 36 | + |
| 37 | + ContextStorageMode mode = args._context_storage; |
| 38 | + if (mode == CTX_STORAGE_OTEL) { |
| 39 | + if (!OtelContexts::initialize()) { |
| 40 | + mode = CTX_STORAGE_PROFILER; |
| 41 | + __atomic_store_n(&_mode, mode, __ATOMIC_RELEASE); |
| 42 | + return false; |
| 43 | + } |
| 44 | + } |
| 45 | + |
| 46 | + __atomic_store_n(&_mode, mode, __ATOMIC_RELEASE); |
| 47 | + __atomic_store_n(&_initialized, true, __ATOMIC_RELEASE); |
| 48 | + return true; |
| 49 | +} |
| 50 | + |
| 51 | +void ContextApi::shutdown() { |
| 52 | + if (!__atomic_load_n(&_initialized, __ATOMIC_ACQUIRE)) { |
| 53 | + return; |
| 54 | + } |
| 55 | + |
| 56 | + if (__atomic_load_n(&_mode, __ATOMIC_ACQUIRE) == CTX_STORAGE_OTEL) { |
| 57 | + OtelContexts::shutdown(); |
| 58 | + } |
| 59 | + |
| 60 | + // Free registered attribute keys |
| 61 | + for (int i = 0; i < _attribute_key_count; i++) { |
| 62 | + free(_attribute_keys[i]); |
| 63 | + _attribute_keys[i] = nullptr; |
| 64 | + } |
| 65 | + _attribute_key_count = 0; |
| 66 | + |
| 67 | + __atomic_store_n(&_initialized, false, __ATOMIC_RELEASE); |
| 68 | +} |
| 69 | + |
| 70 | +bool ContextApi::isInitialized() { |
| 71 | + return __atomic_load_n(&_initialized, __ATOMIC_ACQUIRE); |
| 72 | +} |
| 73 | + |
| 74 | +ContextStorageMode ContextApi::getMode() { |
| 75 | + return __atomic_load_n(&_mode, __ATOMIC_ACQUIRE); |
| 76 | +} |
| 77 | + |
| 78 | +/** |
| 79 | + * Initialize OTel TLS for the current thread on first use. |
| 80 | + * Allocates a per-thread OtelThreadContextRecord and caches it in ProfiledThread. |
| 81 | + * Must be called with signals blocked to prevent musl TLS deadlock. |
| 82 | + */ |
| 83 | +static OtelThreadContextRecord* initializeOtelTls(ProfiledThread* thrd) { |
| 84 | + // Block profiling signals during first TLS access (same pattern as context_tls_v1) |
| 85 | + SignalBlocker blocker; |
| 86 | + |
| 87 | + // Allocate one record per thread — freed in ProfiledThread destructor/releaseFromBuffer |
| 88 | + OtelThreadContextRecord* record = new OtelThreadContextRecord(); |
| 89 | + record->valid = 0; |
| 90 | + record->_reserved = 0; |
| 91 | + record->attrs_data_size = 0; |
| 92 | + |
| 93 | + // First access to custom_labels_current_set_v2 triggers TLS init |
| 94 | + custom_labels_current_set_v2 = nullptr; |
| 95 | + |
| 96 | + // Cache in ProfiledThread for signal-safe cross-thread access |
| 97 | + thrd->markOtelContextInitialized(record); |
| 98 | + |
| 99 | + return record; |
| 100 | +} |
| 101 | + |
| 102 | +void ContextApi::set(u64 span_id, u64 root_span_id) { |
| 103 | + // For Datadog 64-bit trace IDs, trace_id_high is zero. This is valid per OTEP |
| 104 | + // since trace_id_low (rootSpanId) is non-zero — the full 128-bit trace_id is |
| 105 | + // not all-zeros, so the record is considered active. |
| 106 | + setOtel(0, root_span_id, span_id); |
| 107 | +} |
| 108 | + |
| 109 | +void ContextApi::setOtel(u64 trace_id_high, u64 trace_id_low, u64 span_id) { |
| 110 | + ContextStorageMode mode = __atomic_load_n(&_mode, __ATOMIC_ACQUIRE); |
| 111 | + |
| 112 | + if (mode == CTX_STORAGE_OTEL) { |
| 113 | + // All-zero IDs = context detachment per OTEP — use NULL pointer |
| 114 | + if (trace_id_high == 0 && trace_id_low == 0 && span_id == 0) { |
| 115 | + OtelContexts::clear(); |
| 116 | + return; |
| 117 | + } |
| 118 | + |
| 119 | + // Ensure TLS + record are initialized on first use |
| 120 | + ProfiledThread* thrd = ProfiledThread::current(); |
| 121 | + if (thrd == nullptr) return; |
| 122 | + |
| 123 | + if (!thrd->isOtelContextInitialized()) { |
| 124 | + initializeOtelTls(thrd); |
| 125 | + } |
| 126 | + |
| 127 | + OtelContexts::set(trace_id_high, trace_id_low, span_id); |
| 128 | + } else { |
| 129 | + // Profiler mode: use existing TLS |
| 130 | + Context& ctx = Contexts::get(); |
| 131 | + |
| 132 | + __atomic_store_n(&ctx.checksum, 0ULL, __ATOMIC_RELEASE); |
| 133 | + __atomic_store_n(&ctx.spanId, span_id, __ATOMIC_RELAXED); |
| 134 | + __atomic_store_n(&ctx.rootSpanId, trace_id_low, __ATOMIC_RELAXED); |
| 135 | + |
| 136 | + u64 newChecksum = Contexts::checksum(span_id, trace_id_low); |
| 137 | + __atomic_store_n(&ctx.checksum, newChecksum, __ATOMIC_RELEASE); |
| 138 | + } |
| 139 | +} |
| 140 | + |
| 141 | +bool ContextApi::get(u64& span_id, u64& root_span_id) { |
| 142 | + ContextStorageMode mode = __atomic_load_n(&_mode, __ATOMIC_ACQUIRE); |
| 143 | + |
| 144 | + if (mode == CTX_STORAGE_OTEL) { |
| 145 | + u64 trace_high, trace_low; |
| 146 | + if (OtelContexts::get(trace_high, trace_low, span_id)) { |
| 147 | + root_span_id = trace_low; |
| 148 | + return true; |
| 149 | + } |
| 150 | + return false; |
| 151 | + } else { |
| 152 | + Context& ctx = Contexts::get(); |
| 153 | + u64 checksum1 = __atomic_load_n(&ctx.checksum, __ATOMIC_ACQUIRE); |
| 154 | + span_id = __atomic_load_n(&ctx.spanId, __ATOMIC_RELAXED); |
| 155 | + root_span_id = __atomic_load_n(&ctx.rootSpanId, __ATOMIC_RELAXED); |
| 156 | + return checksum1 != 0 && checksum1 == Contexts::checksum(span_id, root_span_id); |
| 157 | + } |
| 158 | +} |
| 159 | + |
| 160 | + |
| 161 | +void ContextApi::clear() { |
| 162 | + ContextStorageMode mode = __atomic_load_n(&_mode, __ATOMIC_ACQUIRE); |
| 163 | + |
| 164 | + if (mode == CTX_STORAGE_OTEL) { |
| 165 | + OtelContexts::clear(); |
| 166 | + } else { |
| 167 | + set(0, 0); |
| 168 | + } |
| 169 | +} |
| 170 | + |
| 171 | +bool ContextApi::setAttribute(uint8_t key_index, const char* value, uint8_t value_len) { |
| 172 | + ContextStorageMode mode = __atomic_load_n(&_mode, __ATOMIC_ACQUIRE); |
| 173 | + |
| 174 | + if (mode == CTX_STORAGE_OTEL) { |
| 175 | + // Ensure TLS + record are initialized on first use |
| 176 | + ProfiledThread* thrd = ProfiledThread::current(); |
| 177 | + if (thrd == nullptr) return false; |
| 178 | + |
| 179 | + if (!thrd->isOtelContextInitialized()) { |
| 180 | + initializeOtelTls(thrd); |
| 181 | + } |
| 182 | + |
| 183 | + // Pre-register the value string in the Dictionary from this JNI thread. |
| 184 | + // The signal handler (writeCurrentContext) will later call bounded_lookup |
| 185 | + // to find the encoding — by pre-registering here, the signal handler |
| 186 | + // only does a read (no malloc), which is async-signal-safe. |
| 187 | + Profiler::instance()->contextValueMap()->bounded_lookup( |
| 188 | + value, value_len, 1 << 16); |
| 189 | + |
| 190 | + return OtelContexts::setAttribute(key_index, value, value_len); |
| 191 | + } else { |
| 192 | + // Profiler mode: register the string and write encoding to Context.tags[] |
| 193 | + if (key_index >= DD_TAGS_CAPACITY) return false; |
| 194 | + |
| 195 | + u32 encoding = Profiler::instance()->contextValueMap()->bounded_lookup( |
| 196 | + value, value_len, 1 << 16); |
| 197 | + if (encoding == INT_MAX) return false; |
| 198 | + |
| 199 | + Context& ctx = Contexts::get(); |
| 200 | + ctx.tags[key_index].value = encoding; |
| 201 | + return true; |
| 202 | + } |
| 203 | +} |
| 204 | + |
| 205 | +void ContextApi::registerAttributeKeys(const char** keys, int count) { |
| 206 | + // Free any previously registered keys |
| 207 | + for (int i = 0; i < _attribute_key_count; i++) { |
| 208 | + free(_attribute_keys[i]); |
| 209 | + _attribute_keys[i] = nullptr; |
| 210 | + } |
| 211 | + |
| 212 | + _attribute_key_count = count < MAX_ATTRIBUTE_KEYS ? count : MAX_ATTRIBUTE_KEYS; |
| 213 | + for (int i = 0; i < _attribute_key_count; i++) { |
| 214 | + _attribute_keys[i] = strdup(keys[i]); |
| 215 | + } |
| 216 | + |
| 217 | + // If in OTEL mode, re-publish process context with thread_ctx_config |
| 218 | + ContextStorageMode mode = __atomic_load_n(&_mode, __ATOMIC_ACQUIRE); |
| 219 | + if (mode == CTX_STORAGE_OTEL) { |
| 220 | + // Build NULL-terminated key array for the process context config |
| 221 | + const char* key_ptrs[MAX_ATTRIBUTE_KEYS + 1]; |
| 222 | + for (int i = 0; i < _attribute_key_count; i++) { |
| 223 | + key_ptrs[i] = _attribute_keys[i]; |
| 224 | + } |
| 225 | + key_ptrs[_attribute_key_count] = nullptr; |
| 226 | + |
| 227 | + otel_thread_ctx_config_data config = { |
| 228 | + .schema_version = "tlsdesc_v1_dev", |
| 229 | + .attribute_key_map = key_ptrs, |
| 230 | + }; |
| 231 | + |
| 232 | +#ifndef OTEL_PROCESS_CTX_NO_READ |
| 233 | + // Re-publish the process context with thread_ctx_config |
| 234 | + // We need to read the current context and re-publish with the config |
| 235 | + otel_process_ctx_read_result read_result = otel_process_ctx_read(); |
| 236 | + if (read_result.success) { |
| 237 | + otel_process_ctx_data data = { |
| 238 | + .deployment_environment_name = read_result.data.deployment_environment_name, |
| 239 | + .service_instance_id = read_result.data.service_instance_id, |
| 240 | + .service_name = read_result.data.service_name, |
| 241 | + .service_version = read_result.data.service_version, |
| 242 | + .telemetry_sdk_language = read_result.data.telemetry_sdk_language, |
| 243 | + .telemetry_sdk_version = read_result.data.telemetry_sdk_version, |
| 244 | + .telemetry_sdk_name = read_result.data.telemetry_sdk_name, |
| 245 | + .resource_attributes = read_result.data.resource_attributes, |
| 246 | + .extra_attributes = read_result.data.extra_attributes, |
| 247 | + .thread_ctx_config = &config, |
| 248 | + }; |
| 249 | + |
| 250 | + otel_process_ctx_publish(&data); |
| 251 | + otel_process_ctx_read_drop(&read_result); |
| 252 | + } |
| 253 | +#endif |
| 254 | + } |
| 255 | +} |
| 256 | + |
| 257 | +const char* ContextApi::getAttributeKey(int index) { |
| 258 | + if (index < 0 || index >= _attribute_key_count) { |
| 259 | + return nullptr; |
| 260 | + } |
| 261 | + return _attribute_keys[index]; |
| 262 | +} |
| 263 | + |
| 264 | +int ContextApi::getAttributeKeyCount() { |
| 265 | + return _attribute_key_count; |
| 266 | +} |
0 commit comments