Generate KnownTags from a tag-conventions registry

buildSrc/build.gradle.kts

@@ -73,6 +73,10 @@ gradlePlugin {
       id = "dd-trace-java.sca-enrichments"
       implementationClass = "datadog.gradle.plugin.sca.ScaEnrichmentsPlugin"
     }
+    create("tag-registry-generator") {
+      id = "dd-trace-java.tag-registry-generator"
+      implementationClass = "datadog.gradle.plugin.tags.TagRegistryGeneratorPlugin"
+    }
   }
 }
 
@@ -102,6 +106,7 @@ dependencies {
   implementation("com.fasterxml.jackson.core:jackson-databind")
   implementation("com.fasterxml.jackson.core:jackson-annotations")
   implementation("com.fasterxml.jackson.core:jackson-core")
+  implementation("com.fasterxml.jackson.dataformat:jackson-dataformat-yaml")
 
   compileOnly(libs.develocity)
 }

buildSrc/src/main/kotlin/datadog/gradle/plugin/tags/GenerateKnownTagsTask.kt

@@ -0,0 +1,38 @@
+package datadog.gradle.plugin.tags
+
+import javax.inject.Inject
+import org.gradle.api.DefaultTask
+import org.gradle.api.file.DirectoryProperty
+import org.gradle.api.file.RegularFileProperty
+import org.gradle.api.model.ObjectFactory
+import org.gradle.api.tasks.CacheableTask
+import org.gradle.api.tasks.InputFile
+import org.gradle.api.tasks.OutputDirectory
+import org.gradle.api.tasks.PathSensitive
+import org.gradle.api.tasks.PathSensitivity
+import org.gradle.api.tasks.TaskAction
+
+/**
+ * Generates the committed tag registry (KnownTags.java + layout reports) from the language-agnostic
+ * {@code tag-conventions.yaml} + the Java overlay. The actual emit lives in [TagRegistryGenerator];
+ * this task just wires the inputs/outputs so Gradle can cache and up-to-date-check it.
+ */
+@CacheableTask
+abstract class GenerateKnownTagsTask @Inject constructor(objects: ObjectFactory) : DefaultTask() {
+  @get:InputFile
+  @get:PathSensitive(PathSensitivity.NONE)
+  val domainYaml: RegularFileProperty = objects.fileProperty()
+
+  @get:InputFile
+  @get:PathSensitive(PathSensitivity.NONE)
+  val overlayYaml: RegularFileProperty = objects.fileProperty()
+
+  @get:OutputDirectory val destinationDirectory: DirectoryProperty = objects.directoryProperty()
+
+  @TaskAction
+  fun generate() {
+    val outDir = destinationDirectory.get().asFile
+    TagRegistryGenerator.generate(domainYaml.get().asFile, overlayYaml.get().asFile, outDir)
+    logger.lifecycle("tag-registry: generated -> $outDir")
+  }
+}

buildSrc/src/main/kotlin/datadog/gradle/plugin/tags/KnownTagsEmitter.kt

@@ -0,0 +1,122 @@
+package datadog.gradle.plugin.tags
+
+/**
+ * Emits the generated `KnownTags.java` from a [TagRegistry]: id constants (literal + a
+ * `// tagId(...)` derivation comment), the `StringIndex.Support` keyOf table, the `globalSerial`
+ * switch `nameOf`, and resolver registration. Mirrors the hand-written KnownTags structure.
+ */
+object KnownTagsEmitter {
+
+  fun emit(reg: TagRegistry, pkg: String, className: String): String {
+    // Sanitize tag names into unique Java constant identifiers.
+    val used = HashSet<String>()
+    val cname = HashMap<String, String>()
+    fun mk(name: String): String {
+      var c = name.uppercase().replace(Regex("[^A-Za-z0-9]"), "_").replace(Regex("_+"), "_").trim('_')
+      if (c.isEmpty() || c[0].isDigit()) c = "T_$c"
+      var u = c
+      var n = 2
+      while (u in used) {
+        u = "${c}_$n"; n++
+      }
+      used.add(u)
+      cname[name] = u
+      return u
+    }
+    reg.virtuals.forEach { mk(it.name) }
+    reg.stored.forEach { mk(it.name) }
+
+    val order = reg.virtuals.map { it.name } + reg.stored.map { it.name } // stable emit order
+    val b = StringBuilder()
+    b.appendLine("package $pkg;")
+    b.appendLine()
+    b.appendLine("import datadog.trace.util.StringIndex;")
+    b.appendLine()
+    b.appendLine("// GENERATED by the tag-registry code generator (dd-trace-java.tag-registry-generator).")
+    b.appendLine("// DO NOT EDIT. Source: tag-conventions.yaml + tag-conventions.java.yaml.")
+    b.appendLine("public final class $className {")
+    b.appendLine("  static final int SLOT_COUNT = ${reg.slotCount};")
+    b.appendLine()
+
+    b.appendLine("  // ---- reserved / virtual (routed to span fields or directives; not stored) ----")
+    for (v in reg.virtuals) {
+      val cn = cname[v.name]!!
+      b.appendLine("  static final int ${cn}_SERIAL = ${v.serial};")
+      b.appendLine("  // tagId(serial=${v.serial}, intercepted=true, slot=NO_SLOT)  [${v.kind}${v.field?.let { " -> $it" } ?: ""}]  \"${v.name}\"")
+      b.appendLine("  public static final long $cn = ${hex(v.id)};")
+    }
+    b.appendLine()
+
+    b.appendLine("  // ---- stored (dense-store slot, or bucketed when slot=NO_SLOT) ----")
+    for (t in reg.stored) {
+      val cn = cname[t.name]!!
+      val slot = if (t.slotted) t.slot.toString() else "NO_SLOT"
+      b.appendLine("  static final int ${cn}_SERIAL = ${t.serial};")
+      b.appendLine("  // tagId(serial=${t.serial}, intercepted=${t.intercepted}, slot=$slot)  <${t.required}${if (t.traceLevel) ", trace-level" else ""}>  \"${t.name}\"")
+      b.appendLine("  public static final long $cn = ${hex(t.id)};")
+    }
+    b.appendLine()
+
+    // keyOf table (open-addressed, via StringIndex.Support).
+    b.appendLine("  private static final String[] KEYOF_NAMES = {")
+    order.forEach { b.appendLine("    \"$it\",") }
+    b.appendLine("  };")
+    b.appendLine("  private static final long[] KEYOF_VALUES = {")
+    order.forEach { b.appendLine("    ${cname[it]},") }
+    b.appendLine("  };")
+    b.appendLine("  private static final int[] KEYOF_HASHES;")
+    b.appendLine("  private static final String[] KEYOF_KEYS;")
+    b.appendLine("  private static final long[] KEYOF_IDS;")
+    b.appendLine("  static {")
+    b.appendLine("    StringIndex.Data data = StringIndex.Support.create(KEYOF_NAMES);")
+    b.appendLine("    long[] ids = new long[data.names.length];")
+    b.appendLine("    for (int j = 0; j < KEYOF_NAMES.length; j++) {")
+    b.appendLine("      ids[StringIndex.Support.indexOf(data.hashes, data.names, KEYOF_NAMES[j])] = KEYOF_VALUES[j];")
+    b.appendLine("    }")
+    b.appendLine("    KEYOF_HASHES = data.hashes;")
+    b.appendLine("    KEYOF_KEYS = data.names;")
+    b.appendLine("    KEYOF_IDS = ids;")
+    b.appendLine("  }")
+    b.appendLine()
+
+    // Resolver.
+    b.appendLine("  static final KnownTagCodec.Resolver RESOLVER =")
+    b.appendLine("      new KnownTagCodec.Resolver() {")
+    b.appendLine("        @Override")
+    b.appendLine("        public String nameOf(long tagId) {")
+    b.appendLine("          switch (KnownTagCodec.globalSerial(tagId)) {")
+    for (name in order) {
+      b.appendLine("            case ${cname[name]}_SERIAL:")
+      b.appendLine("              return \"$name\";")
+    }
+    b.appendLine("            default:")
+    b.appendLine("              return null;")
+    b.appendLine("          }")
+    b.appendLine("        }")
+    b.appendLine()
+    b.appendLine("        @Override")
+    b.appendLine("        public int slotCount() {")
+    b.appendLine("          return SLOT_COUNT;")
+    b.appendLine("        }")
+    b.appendLine()
+    b.appendLine("        @Override")
+    b.appendLine("        public long keyOf(String name) {")
+    b.appendLine("          int slot = StringIndex.Support.indexOf(KEYOF_HASHES, KEYOF_KEYS, name);")
+    b.appendLine("          return slot < 0 ? 0L : KEYOF_IDS[slot];")
+    b.appendLine("        }")
+    b.appendLine("      };")
+    b.appendLine()
+    b.appendLine("  static {")
+    b.appendLine("    KnownTagCodec.register(RESOLVER);")
+    b.appendLine("  }")
+    b.appendLine()
+    b.appendLine("  /** Forces resolver registration by triggering <clinit>. Idempotent. */")
+    b.appendLine("  public static void init() {}")
+    b.appendLine()
+    b.appendLine("  private $className() {}")
+    b.appendLine("}")
+    return b.toString()
+  }
+
+  private fun hex(id: Long): String = "0x%016XL".format(id)
+}

buildSrc/src/main/kotlin/datadog/gradle/plugin/tags/TagConventions.kt

@@ -0,0 +1,150 @@
+package datadog.gradle.plugin.tags
+
+/**
+ * Parsed tag-conventions domain model + the per-type tag-set resolver. Language-agnostic: it knows
+ * only structure (extends / include / applies) and per-tag semantics (name / type / required /
+ * source). Id assignment, coloring, and emission are layered on top of the resolved sets.
+ */
+class TagConventions
+private constructor(
+  private val spanTypes: Map<String, SpanType>,
+  private val mixins: Map<String, Mixin>,
+  private val traceLevel: List<Tag>,
+) {
+  /** A tag declaration (domain semantics only). */
+  data class Tag(
+    val name: String,
+    val type: String,
+    val required: String,
+  )
+
+  data class SpanType(
+    val name: String,
+    val abstract: Boolean,
+    val extends: String?,
+    val include: List<String>,
+    val tags: List<Tag>,
+  )
+
+  data class Mixin(
+    val name: String,
+    val appliesAll: Boolean,
+    val appliesTo: Set<String>,
+    val tags: List<Tag>,
+  )
+
+  /** Concrete (instantiable) span types — the ones a layout is computed for. */
+  fun concreteTypes(): List<String> =
+    spanTypes.values.filter { !it.abstract }.map { it.name }.sorted()
+
+  /**
+   * resolved(type) = own tags + tags up the `extends` chain (incl. base) + tags of every mixin the
+   * type or an ancestor `include`s + tags of every mixin whose `applies` matches. De-duped by tag
+   * name (first occurrence wins). Base-first order, so it is stable across runs.
+   */
+  fun resolve(typeName: String): List<Tag> {
+    val result = LinkedHashMap<String, Tag>()
+    fun add(t: Tag) = result.putIfAbsent(t.name, t)
+
+    val chain = ArrayList<SpanType>()
+    var cur: SpanType? = spanTypes[typeName]
+    while (cur != null) {
+      chain.add(cur)
+      cur = cur.extends?.let { spanTypes[it] }
+    }
+    for (st in chain.asReversed()) {
+      st.tags.forEach { add(it) }
+      for (mixinName in st.include) mixins[mixinName]?.tags?.forEach { add(it) }
+    }
+    val chainNames = chain.map { it.name }.toSet()
+    for (mx in mixins.values) {
+      if (mx.appliesAll || mx.appliesTo.any { it in chainNames }) mx.tags.forEach { add(it) }
+    }
+    return result.values.toList()
+  }
+
+  /** The explicit trace-level tier tags (their own TagMap "type" on the TraceSegment). */
+  fun traceLevelTags(): List<Tag> = traceLevel
+
+  /** Full stored-tag universe (concrete span types' resolves + trace-level), de-duped by name. */
+  fun allStoredTags(): List<Tag> {
+    val union = LinkedHashMap<String, Tag>()
+    for (type in concreteTypes()) for (t in resolve(type)) union.putIfAbsent(t.name, t)
+    for (t in traceLevel) union.putIfAbsent(t.name, t)
+    return union.values.toList()
+  }
+
+  /** Resolved tag NAMES per concrete type — the co-occurrence structure the colorer needs. */
+  fun coOccurrence(): Map<String, Set<String>> =
+    concreteTypes().associateWith { resolve(it).map { t -> t.name }.toSet() }
+
+  /**
+   * Full composition for a type as (origin, tag) pairs, in composition order and NOT de-duped, so a
+   * tag contributed by more than one source shows up more than once. Origin is the contributing
+   * span type (via extends), `incl:<mixin>` (via include), or `appl:<mixin>` (via applies).
+   */
+  fun compose(typeName: String): List<Pair<String, Tag>> {
+    val out = ArrayList<Pair<String, Tag>>()
+    val chain = ArrayList<SpanType>()
+    var cur: SpanType? = spanTypes[typeName]
+    while (cur != null) {
+      chain.add(cur)
+      cur = cur.extends?.let { spanTypes[it] }
+    }
+    for (st in chain.asReversed()) {
+      st.tags.forEach { out.add(st.name to it) }
+      for (mixinName in st.include) mixins[mixinName]?.tags?.forEach { out.add("incl:$mixinName" to it) }
+    }
+    val chainNames = chain.map { it.name }.toSet()
+    for (mx in mixins.values) {
+      if (mx.appliesAll || mx.appliesTo.any { it in chainNames }) {
+        mx.tags.forEach { out.add("appl:${mx.name}" to it) }
+      }
+    }
+    return out
+  }
+
+  companion object {
+    @Suppress("UNCHECKED_CAST")
+    fun parse(root: Map<String, Any?>): TagConventions {
+      val spanTypesRaw = (root["span_types"] as? Map<String, Any?>) ?: emptyMap()
+      val spanTypes =
+        spanTypesRaw.mapValues { (name, v) ->
+          val m = v as Map<String, Any?>
+          SpanType(
+            name = name,
+            abstract = (m["abstract"] as? Boolean) ?: false,
+            extends = m["extends"] as? String,
+            include = (m["include"] as? List<String>) ?: emptyList(),
+            tags = tagList(m["tags"]),
+          )
+        }
+
+      val mixinsRaw = (root["mixins"] as? Map<String, Any?>) ?: emptyMap()
+      val mixins =
+        mixinsRaw.mapValues { (name, v) ->
+          val m = v as Map<String, Any?>
+          val applies = m["applies"]
+          Mixin(
+            name = name,
+            appliesAll = applies == "all",
+            appliesTo = if (applies is List<*>) applies.map { it.toString() }.toSet() else emptySet(),
+            tags = tagList(m["tags"]),
+          )
+        }
+
+      val traceLevel = tagList((root["trace_level"] as? Map<String, Any?>)?.get("tags"))
+      return TagConventions(spanTypes, mixins, traceLevel)
+    }
+
+    @Suppress("UNCHECKED_CAST")
+    private fun tagList(tags: Any?): List<Tag> =
+      (tags as? List<Map<String, Any?>>)?.map { m ->
+        Tag(
+          name = m["tag"].toString(),
+          type = (m["type"] as? String) ?: "string",
+          required = (m["required"] as? String) ?: "optional",
+        )
+      } ?: emptyList()
+  }
+}