diff --git a/.github/workflows/continuous.yml b/.github/workflows/continuous.yml index a3733af34..6a5c83868 100644 --- a/.github/workflows/continuous.yml +++ b/.github/workflows/continuous.yml @@ -26,7 +26,7 @@ jobs: - name: Set up Java uses: actions/setup-java@v5 with: - java-version: 11 + java-version: 21 distribution: temurin - name: Checkout uses: actions/checkout@v6 diff --git a/btrace-client/src/main/java/io/btrace/client/Client.java b/btrace-client/src/main/java/io/btrace/client/Client.java index 636c23f1a..da948bda0 100644 --- a/btrace-client/src/main/java/io/btrace/client/Client.java +++ b/btrace-client/src/main/java/io/btrace/client/Client.java @@ -801,7 +801,7 @@ public void attach(String pid, String agentPath, String sysCp, String bootCp) th } } - void connectAndListProbes(String host, CommandListener listener) throws IOException { + public void connectAndListProbes(String host, CommandListener listener) throws IOException { if (sock != null) { throw new IllegalStateException(); } @@ -1111,11 +1111,11 @@ public synchronized void close() throws IOException { reset(); } - boolean isDisconnected() { + public boolean isDisconnected() { return disconnected; } - void disconnect() throws IOException { + public void disconnect() throws IOException { disconnected = true; if (log.isDebugEnabled()) { log.debug("sending DISCONNECT request to agent"); diff --git a/btrace-client/src/main/java/io/btrace/extcli/Installer.java b/btrace-client/src/main/java/io/btrace/extcli/Installer.java index b1d3e414f..07923741b 100644 --- a/btrace-client/src/main/java/io/btrace/extcli/Installer.java +++ b/btrace-client/src/main/java/io/btrace/extcli/Installer.java @@ -108,7 +108,9 @@ static void install(String target, List repos, String id, boolean dryRun throw new IOException("Failed to download extension from provided repositories."); } else { throw new IllegalArgumentException( - "Unrecognized input: provide a zip path, URL, or group:artifact:version"); + "Unrecognised target: '" + + target + + "'. Expected a URL, .zip path, or groupId:artifactId:version coordinate."); } // Validate zip contains -api.jar and -impl.jar, and install diff --git a/btrace-client/src/test/java/io/btrace/extcli/ExtensionListerTest.java b/btrace-client/src/test/java/io/btrace/extcli/ExtensionListerTest.java index cb634139e..683869587 100644 --- a/btrace-client/src/test/java/io/btrace/extcli/ExtensionListerTest.java +++ b/btrace-client/src/test/java/io/btrace/extcli/ExtensionListerTest.java @@ -36,13 +36,11 @@ class ExtensionListerTest { private final ByteArrayOutputStream errContent = new ByteArrayOutputStream(); private final PrintStream originalOut = System.out; private final PrintStream originalErr = System.err; - private String originalBtraceHome; @BeforeEach void setUpStreams() { System.setOut(new PrintStream(outContent)); System.setErr(new PrintStream(errContent)); - originalBtraceHome = System.getenv("BTRACE_HOME"); } @AfterEach @@ -68,17 +66,6 @@ void listFromBtraceHome() throws IOException { assertNotNull(output); } - @Test - void listWithJsonFormat() throws IOException { - ExtensionLister.list(true); - - String output = outContent.toString(); - // Should output valid JSON (starts with [ and ends with ]) - assertTrue( - output.trim().startsWith("[") && output.trim().endsWith("]"), - "JSON output should be an array"); - } - @Test void listHandlesEmptyDirectories() throws IOException { // Create empty extensions directory diff --git a/btrace-compiler/src/main/java/io/btrace/compiler/VerifierVisitor.java b/btrace-compiler/src/main/java/io/btrace/compiler/VerifierVisitor.java index 5a3fcf45c..f49e22343 100644 --- a/btrace-compiler/src/main/java/io/btrace/compiler/VerifierVisitor.java +++ b/btrace-compiler/src/main/java/io/btrace/compiler/VerifierVisitor.java @@ -537,14 +537,22 @@ public Void visitVariable(VariableTree vt, Void p) { /** * Returns true if the given service class name is declared by any extension. * - *

Compile-time validation notes: - This check operates without loading classes. It inspects - * extension metadata (BTrace-Extension-Services in MANIFEST.MF and legacy - * META-INF/btrace-extension.properties) to verify that an @Injected service type is declared by - * some extension. - It complements the bytecode-time check in instr (BTraceProbeNode) and the - * runtime reflection-based validation in the agent (Client#validateDeclaredServices). The latter - * ensures correctness under the actual runtime classloader/JPMS environment. + *

Compile-time validation notes: checks first via the annotation processing type model (which + * sees -cp JARs), then falls back to scanning jar manifests on the JVM classpath. The agent + * performs a definitive runtime check (Client#validateDeclaredServices). */ private boolean isDeclaredExtensionService(String serviceClassName) { + // Primary: use the annotation processing Elements API — sees the compilation classpath + TypeElement te = verifier.getElementUtils().getTypeElement(serviceClassName); + if (te != null) { + for (javax.lang.model.element.AnnotationMirror am : te.getAnnotationMirrors()) { + String annName = am.getAnnotationType().asElement().toString(); + if ("io.btrace.core.extensions.ServiceDescriptor".equals(annName)) { + return true; + } + } + } + // Fallback: scan JAR manifests visible to the JVM classloader String resourceName = serviceClassName.replace('.', '/') + ".class"; ClassLoader cl = Thread.currentThread().getContextClassLoader(); if (cl == null) { @@ -566,7 +574,6 @@ private boolean isDeclaredExtensionService(String serviceClassName) { } } } - // Fallback: scan system classpath jars String cp = System.getProperty("java.class.path", ""); String[] parts = cp.split(java.io.File.pathSeparator); for (String p : parts) { diff --git a/btrace-dist/build.gradle b/btrace-dist/build.gradle index c014c44e1..b49c4237d 100644 --- a/btrace-dist/build.gradle +++ b/btrace-dist/build.gradle @@ -638,7 +638,7 @@ task buildDockerContext(type: Copy, dependsOn: [btraceJar, copyDtraceLib, proces from "${distTarget}" into layout.buildDirectory.dir("docker-context/btrace") - fileMode = 0644 + filePermissions { unix(0644) } doLast { // Copy entrypoint script @@ -807,7 +807,7 @@ task buildDockerImages { description "Generates Javadoc API documentation for the btrace-${name}." title = "btrace-${name}" - destinationDir = file("${project.docsDir}/${name}/javadoc") + destinationDir = layout.buildDirectory.dir("docs/${name}/javadoc").get().asFile classpath = files(compileJava.destinationDirectory) + configurations.artifact.asFileTree options.addStringOption('Xdoclint:all,-missing', '-quiet') failOnError false diff --git a/btrace-dist/src/main/resources/samples/ContractCheck.java b/btrace-dist/src/main/resources/samples/ContractCheck.java new file mode 100644 index 000000000..bc94f2ffc --- /dev/null +++ b/btrace-dist/src/main/resources/samples/ContractCheck.java @@ -0,0 +1,103 @@ +import io.btrace.contracts.ContractService; +import io.btrace.core.annotations.BTrace; +import io.btrace.core.annotations.Duration; +import io.btrace.core.annotations.Injected; +import io.btrace.core.annotations.Kind; +import io.btrace.core.annotations.Location; +import io.btrace.core.annotations.OnEvent; +import io.btrace.core.annotations.OnMethod; +import io.btrace.core.annotations.OnTimer; +import io.btrace.core.annotations.ProbeClassName; +import io.btrace.core.annotations.ProbeMethodName; +import io.btrace.core.annotations.Return; + +import static io.btrace.core.BTraceUtils.*; + +/** + * Runtime behavioral contracts: latency budgets, call-rate limits, null-safety, and tagged + * code path profiling — without modifying the target code. + * + *

Requires the btrace-contracts extension in $BTRACE_HOME/extensions/. + * + *

Attach to a running JVM: + * 

+ * btrace <pid> ContractCheck.java
+ *
+ * + *

Trigger an on-demand summary: + * 

+ * btrace <pid> --event summary
+ *
+ */ +@BTrace +public class ContractCheck { + + @Injected + private static ContractService contracts; + + // ==================== Latency budgets ==================== + + /** + * Enforce a 500ms latency budget on all methods in the target package. + * Adjust the clazz pattern to match your project structure. + */ + @OnMethod( + clazz = "/com\\.myapp\\..*/", + method = "/.*/", + location = @Location(Kind.RETURN)) + public static void checkLatency( + @ProbeClassName String cls, + @ProbeMethodName String method, + @Duration long dur) { + contracts.checkLatency(strcat(cls, strcat(".", method)), dur, 500_000_000L); + } + + // ==================== Null safety ==================== + + @OnMethod( + clazz = "/com\\.myapp\\..*/", + method = "/.*/", + location = @Location(Kind.RETURN)) + public static void checkNullReturn( + @ProbeClassName String cls, + @ProbeMethodName String method, + @Return Object ret) { + contracts.checkNotNull(strcat(cls, strcat(".", method)), ret); + } + + // ==================== Tagged path comparison ==================== + + // Example: compare two implementations of the same operation. + // Tag them differently so getSummary() renders them side by side. + + @OnMethod( + clazz = "com.myapp.CachedQueryService", + method = "query", + location = @Location(Kind.RETURN)) + public static void onCachedQuery(@Duration long dur) { + contracts.trackCodePath("QueryService.query", dur, "cached"); + } + + @OnMethod( + clazz = "com.myapp.DirectQueryService", + method = "query", + location = @Location(Kind.RETURN)) + public static void onDirectQuery(@Duration long dur) { + contracts.trackCodePath("QueryService.query", dur, "direct"); + } + + // ==================== Reporting ==================== + + @OnTimer(10000) + public static void periodicCheck() { + if (contracts.hasViolations()) { + println("=== CONTRACT VIOLATIONS DETECTED ==="); + println(contracts.getSummary()); + } + } + + @OnEvent("summary") + public static void onDemandSummary() { + println(contracts.getSummary()); + } +} diff --git a/btrace-dist/src/main/resources/samples/GpuBridge.java b/btrace-dist/src/main/resources/samples/GpuBridge.java new file mode 100644 index 000000000..b0c81ee52 --- /dev/null +++ b/btrace-dist/src/main/resources/samples/GpuBridge.java @@ -0,0 +1,85 @@ +import io.btrace.core.annotations.BTrace; +import io.btrace.core.annotations.Duration; +import io.btrace.core.annotations.Injected; +import io.btrace.core.annotations.Kind; +import io.btrace.core.annotations.Location; +import io.btrace.core.annotations.OnEvent; +import io.btrace.core.annotations.OnMethod; +import io.btrace.core.annotations.OnTimer; +import io.btrace.gpu.GpuBridgeService; + +import static io.btrace.core.BTraceUtils.*; + +/** + * Traces GPU model inference via ONNX Runtime and DJL (Deep Java Library). + * Tracks inference latency, batch sizes, and model load times. + * + *

Attach to a JVM running ONNX or DJL inference: + * 

+ * btrace <pid> GpuBridge.java
+ *
+ */ +@BTrace +public class GpuBridge { + + @Injected + private static GpuBridgeService gpu; + + // ==================== ONNX Runtime ==================== + + @OnMethod( + clazz = "ai.onnxruntime.OrtSession", + method = "run", + location = @Location(Kind.RETURN)) + public static void onOnnxInference(@Duration long dur) { + gpu.recordInference("onnx", "session", dur); + } + + @OnMethod( + clazz = "ai.onnxruntime.OrtSession", + method = "", + location = @Location(Kind.RETURN)) + public static void onOnnxModelLoad(@Duration long dur) { + gpu.recordModelLoad("onnx", "session", dur); + } + + // ==================== DJL (Deep Java Library) ==================== + + @OnMethod( + clazz = "/ai\\.djl\\.inference\\.Predictor/", + method = "predict", + location = @Location(Kind.RETURN)) + public static void onDjlPredict(@Duration long dur) { + gpu.recordInference("djl", "predictor", dur); + } + + @OnMethod( + clazz = "/ai\\.djl\\.repository\\.zoo\\.ModelZoo/", + method = "loadModel", + location = @Location(Kind.RETURN)) + public static void onDjlModelLoad(@Duration long dur) { + gpu.recordModelLoad("djl", "model-zoo", dur); + } + + // ==================== TensorFlow Java ==================== + + @OnMethod( + clazz = "/org\\.tensorflow\\.Session/", + method = "run", + location = @Location(Kind.RETURN)) + public static void onTensorFlowRun(@Duration long dur) { + gpu.recordInference("tensorflow", "session", dur); + } + + // ==================== Periodic summary ==================== + + @OnTimer(30000) + public static void periodicSummary() { + println(gpu.getSummary()); + } + + @OnEvent("summary") + public static void onDemandSummary() { + println(gpu.getSummary()); + } +} diff --git a/btrace-dist/src/main/resources/samples/LlmTrace.java b/btrace-dist/src/main/resources/samples/LlmTrace.java new file mode 100644 index 000000000..941d3ddbd --- /dev/null +++ b/btrace-dist/src/main/resources/samples/LlmTrace.java @@ -0,0 +1,119 @@ +/* + * Copyright (c) 2024, Oracle and/or its affiliates. All rights reserved. + * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER. + * + * This code is free software; you can redistribute it and/or modify it + * under the terms of the GNU General Public License version 2 only, as + * published by the Free Software Foundation. Oracle designates this + * particular file as subject to the Classpath exception as provided + * by Oracle in the LICENSE file that accompanied this code. + * + * This code 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 + * version 2 for more details (a copy is included in the LICENSE file that + * accompanied this code). + * + * You should have received a copy of the GNU General Public License version + * 2 along with this work; if not, write to the Free Software Foundation, + * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA. + * + * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA + * or visit www.oracle.com if you need additional information or have any + * questions. + */ + +import io.btrace.core.annotations.BTrace; +import io.btrace.core.annotations.Duration; +import io.btrace.core.annotations.Injected; +import io.btrace.core.annotations.Kind; +import io.btrace.core.annotations.Location; +import io.btrace.core.annotations.OnEvent; +import io.btrace.core.annotations.OnMethod; +import io.btrace.core.annotations.OnTimer; +import io.btrace.core.annotations.ProbeClassName; +import io.btrace.llm.LlmTraceService; + +import static io.btrace.core.BTraceUtils.*; + +/** + * Sample BTrace script that traces LLM API calls using the btrace-llm-trace extension. + * + *

This is a generic template. It intercepts Langchain4j's ChatLanguageModel.generate() + * method and records call metrics. Adapt the @OnMethod annotations to target your specific + * LLM SDK (OpenAI Java SDK, Anthropic Java SDK, Spring AI, etc.). + * + *

Usage: + * btrace <PID> LlmTrace.java + * + *

Send a named event "summary" to print the current stats: + * (Ctrl-C, option 3, enter "summary") + */ +@BTrace +public class LlmTrace { + + @Injected + private static LlmTraceService llm; + + /** + * Trace Langchain4j ChatLanguageModel.generate() calls. + * Captures latency on every call completion. + */ + @OnMethod( + clazz = "+dev.langchain4j.model.chat.ChatLanguageModel", + method = "generate", + location = @Location(Kind.RETURN)) + public static void onLangchain4jGenerate( + @ProbeClassName String className, + @Duration long duration) { + // Model name extracted from the class; token counts need return value parsing + // For a production script, parse the Response return value + llm.recordCall(className, duration); + println(strcat(strcat("LLM call: ", className), strcat(" ", strcat(str(duration / 1000000L), "ms")))); + } + + /** + * Trace Langchain4j StreamingChatLanguageModel calls. + */ + @OnMethod( + clazz = "+dev.langchain4j.model.chat.StreamingChatLanguageModel", + method = "generate", + location = @Location(Kind.RETURN)) + public static void onLangchain4jStreaming( + @ProbeClassName String className, + @Duration long duration) { + llm.call(className).streaming().duration(duration).record(); + } + + /** + * Trace errors from any ChatLanguageModel implementation. + */ + @OnMethod( + clazz = "+dev.langchain4j.model.chat.ChatLanguageModel", + method = "generate", + location = @Location(Kind.ERROR)) + public static void onLangchain4jError( + @ProbeClassName String className, + @Duration long duration) { + llm.recordError(className, "exception", duration); + println(strcat("LLM ERROR in: ", className)); + } + + /** + * Print summary on named event "summary". + */ + @OnEvent("summary") + public static void onSummary() { + println(llm.getSummary()); + } + + /** + * Print summary periodically (every 30 seconds). + */ + @OnTimer(30000) + public static void onTimer() { + if (llm.getTotalCalls() > 0) { + println(llm.getSummary()); + } + } +} diff --git a/btrace-dist/src/main/resources/samples/RagQuality.java b/btrace-dist/src/main/resources/samples/RagQuality.java new file mode 100644 index 000000000..960fee8b8 --- /dev/null +++ b/btrace-dist/src/main/resources/samples/RagQuality.java @@ -0,0 +1,79 @@ +import io.btrace.core.annotations.BTrace; +import io.btrace.core.annotations.Duration; +import io.btrace.core.annotations.Injected; +import io.btrace.core.annotations.Kind; +import io.btrace.core.annotations.Location; +import io.btrace.core.annotations.OnEvent; +import io.btrace.core.annotations.OnMethod; +import io.btrace.core.annotations.OnTimer; +import io.btrace.rag.RagQualityService; + +import static io.btrace.core.BTraceUtils.*; + +/** + * Traces RAG pipeline performance: vector DB query latency, similarity + * scores, and empty retrieval rates. Targets Pinecone Java client. + * + *

Attach to a JVM running a RAG pipeline: + * 

+ * btrace <pid> RagQuality.java
+ *
+ */ +@BTrace +public class RagQuality { + + @Injected + private static RagQualityService rag; + + // ==================== Pinecone ==================== + + @OnMethod( + clazz = "/io\\.pinecone\\..*/", + method = "query", + location = @Location(Kind.RETURN)) + public static void onPineconeQuery(@Duration long dur) { + rag.recordQuery("pinecone", dur); + } + + // ==================== Milvus ==================== + + @OnMethod( + clazz = "/io\\.milvus\\.client\\..*/", + method = "search", + location = @Location(Kind.RETURN)) + public static void onMilvusSearch(@Duration long dur) { + rag.recordQuery("milvus", dur); + } + + // ==================== Weaviate ==================== + + @OnMethod( + clazz = "/io\\.weaviate\\.client\\..*/", + method = "/get|search/", + location = @Location(Kind.RETURN)) + public static void onWeaviateQuery(@Duration long dur) { + rag.recordQuery("weaviate", dur); + } + + // ==================== Chroma ==================== + + @OnMethod( + clazz = "/tech\\.amikos\\.chromadb\\..*/", + method = "query", + location = @Location(Kind.RETURN)) + public static void onChromaQuery(@Duration long dur) { + rag.recordQuery("chroma", dur); + } + + // ==================== Periodic summary ==================== + + @OnTimer(30000) + public static void periodicSummary() { + println(rag.getSummary()); + } + + @OnEvent("summary") + public static void onDemandSummary() { + println(rag.getSummary()); + } +} diff --git a/btrace-extensions/btrace-contracts/build.gradle b/btrace-extensions/btrace-contracts/build.gradle new file mode 100644 index 000000000..9ec12efda --- /dev/null +++ b/btrace-extensions/btrace-contracts/build.gradle @@ -0,0 +1,35 @@ +plugins { + id 'io.btrace.extension' + alias(libs.plugins.shadow) +} + +java { + sourceCompatibility = 8 + targetCompatibility = 8 +} + +compileJava { + javaCompiler = javaToolchains.compilerFor { + languageVersion.set(JavaLanguageVersion.of(11)) + } +} + +btraceExtension { + id = 'btrace-contracts' + name = 'BTrace Contracts' + description = 'Runtime behavioral contracts: latency budgets, call-rate limits, value assertions, and tagged code path profiling' + services = ['io.btrace.contracts.ContractService'] +} + +dependencies { + apiCompileOnly 'org.jetbrains:annotations:26.1.0' + apiCompileOnly project(':btrace-core') + implCompileOnly 'org.jetbrains:annotations:26.1.0' + implImplementation project(':btrace-core') + + testImplementation libs.junit.jupiter +} + +test { + useJUnitPlatform() +} diff --git a/btrace-extensions/btrace-contracts/src/main/java/io/btrace/contracts/ContractService.java b/btrace-extensions/btrace-contracts/src/main/java/io/btrace/contracts/ContractService.java new file mode 100644 index 000000000..3019e4ee4 --- /dev/null +++ b/btrace-extensions/btrace-contracts/src/main/java/io/btrace/contracts/ContractService.java @@ -0,0 +1,99 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.contracts; + +import io.btrace.core.extensions.ServiceDescriptor; +import org.jetbrains.annotations.Nullable; + +/** + * BTrace extension for enforcing runtime behavioral contracts. + * + *

Validates that methods respect declared invariants at runtime — latency budgets, call-rate + * limits, null-safety, and value ranges — without modifying the target code. Violations are counted + * and reported; the target application is never interrupted. + * + *

Optionally tracks call count and average latency per user-supplied tag, so you can compare any + * two code paths (e.g. cached vs uncached, v1 vs v2, impl A vs impl B) side by side. + * + *

Usage in a BTrace script: + * + * 

+ * @Injected ContractService contracts;
+ *
+ * @OnMethod(clazz = "com.app.Service", method = "query",
+ *           location = @Location(Kind.RETURN))
+ * void onReturn(@Duration long dur) {
+ *     contracts.checkLatency("Service.query", dur, 200_000_000L); // 200ms budget
+ * }
+ *
+ */ +@ServiceDescriptor +public interface ContractService { + + // ==================== Contract checks ==================== + + /** + * Checks that a method's latency does not exceed the budget. Records a violation if {@code + * durationNanos > budgetNanos}. + */ + void checkLatency(String contract, long durationNanos, long budgetNanos); + + /** + * Checks that the call rate does not exceed the limit per second. Uses a sliding window to detect + * bursts. + */ + void checkCallRate(String contract, int maxPerSecond); + + /** Asserts a boolean condition. Records a violation if false. */ + void assertCondition(String contract, boolean condition, String message); + + /** Checks that a numeric value is within bounds (inclusive). Records a violation if not. */ + void checkRange(String contract, long value, long min, long max); + + /** Checks that a return value is not null. Records a violation if null. */ + void checkNotNull(String contract, Object value); + + // ==================== Tracking ==================== + + /** + * Records a call against the given tag for the contract. Call count and cumulative duration are + * accumulated per tag; {@link #getSummary()} renders all tags side by side. The tag is + * user-defined — common examples: {@code "cached"}/{@code "uncached"}, {@code "v1"}/{@code "v2"}, + * {@code "fast-path"}/{@code "slow-path"}. + */ + void trackCodePath(String contract, long durationNanos, String tag); + + // ==================== Reporting ==================== + + /** Returns a formatted summary of all contract checks, violations, and tracked paths. */ + @Nullable String getSummary(); + + /** Returns total number of contract violations across all contracts. */ + long getTotalViolations(); + + /** Returns number of violations for a specific contract. */ + long getViolations(String contract); + + /** Returns total number of contract checks performed. */ + long getTotalChecks(); + + /** Returns true if any contract has been violated. */ + boolean hasViolations(); + + /** Resets all metrics and violation history. */ + void reset(); +} diff --git a/btrace-extensions/btrace-contracts/src/main/java/io/btrace/contracts/ContractServiceImpl.java b/btrace-extensions/btrace-contracts/src/main/java/io/btrace/contracts/ContractServiceImpl.java new file mode 100644 index 000000000..c5d9392a6 --- /dev/null +++ b/btrace-extensions/btrace-contracts/src/main/java/io/btrace/contracts/ContractServiceImpl.java @@ -0,0 +1,343 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.contracts; + +import io.btrace.core.extensions.Extension; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicLong; + +/** Thread-safe behavioral contract enforcement with lock-free statistics. */ +public final class ContractServiceImpl extends Extension implements ContractService { + + private final ConcurrentHashMap contracts = new ConcurrentHashMap<>(); + private final ConcurrentHashMap rateWindows = new ConcurrentHashMap<>(); + // contract -> tag -> stats + private final ConcurrentHashMap> taggedPaths = + new ConcurrentHashMap<>(); + + // ==================== Contract checks ==================== + + @Override + public void checkLatency(String contract, long durationNanos, long budgetNanos) { + ContractStats stats = getOrCreate(contract); + stats.checks.incrementAndGet(); + if (durationNanos > budgetNanos) { + stats.violations.incrementAndGet(); + stats.lastViolationMessage = + "Latency " + + (durationNanos / 1_000_000) + + "ms exceeded budget " + + (budgetNanos / 1_000_000) + + "ms"; + } + stats.totalDurationNanos.addAndGet(durationNanos); + updateMinMax(stats.minDurationNanos, stats.maxDurationNanos, durationNanos); + } + + @Override + public void checkCallRate(String contract, int maxPerSecond) { + ContractStats stats = getOrCreate(contract); + stats.checks.incrementAndGet(); + + RateWindow window = rateWindows.computeIfAbsent(contract, k -> new RateWindow()); + long now = System.nanoTime(); + long count = window.recordAndCount(now); + if (count > maxPerSecond) { + stats.violations.incrementAndGet(); + stats.lastViolationMessage = "Rate " + count + "/sec exceeded limit " + maxPerSecond + "/sec"; + } + } + + @Override + public void assertCondition(String contract, boolean condition, String message) { + ContractStats stats = getOrCreate(contract); + stats.checks.incrementAndGet(); + if (!condition) { + stats.violations.incrementAndGet(); + stats.lastViolationMessage = message; + } + } + + @Override + public void checkRange(String contract, long value, long min, long max) { + ContractStats stats = getOrCreate(contract); + stats.checks.incrementAndGet(); + if (value < min || value > max) { + stats.violations.incrementAndGet(); + stats.lastViolationMessage = "Value " + value + " outside range [" + min + ", " + max + "]"; + } + } + + @Override + public void checkNotNull(String contract, Object value) { + ContractStats stats = getOrCreate(contract); + stats.checks.incrementAndGet(); + if (value == null) { + stats.violations.incrementAndGet(); + stats.lastViolationMessage = "Unexpected null return"; + } + } + + // ==================== Tracking ==================== + + @Override + public void trackCodePath(String contract, long durationNanos, String tag) { + ConcurrentHashMap byTag = + taggedPaths.computeIfAbsent(contract, k -> new ConcurrentHashMap<>()); + CodePathStats s = byTag.computeIfAbsent(tag, k -> new CodePathStats()); + s.calls.incrementAndGet(); + s.totalDurationNanos.addAndGet(durationNanos); + } + + // ==================== Reporting ==================== + + @Override + public String getSummary() { + if (contracts.isEmpty() && taggedPaths.isEmpty()) { + return "No contracts checked."; + } + + StringBuilder sb = new StringBuilder(); + sb.append("=== Contract Summary ===\n\n"); + + long totalChecks = 0; + long totalViolations = 0; + + for (Map.Entry entry : contracts.entrySet()) { + String name = entry.getKey(); + ContractStats s = entry.getValue(); + long checks = s.checks.get(); + long violations = s.violations.get(); + totalChecks += checks; + totalViolations += violations; + + sb.append("Contract: ").append(name).append("\n"); + sb.append(" Checks: ").append(checks); + if (violations > 0) { + sb.append(" | VIOLATIONS: ").append(violations); + sb.append(" (").append(violations * 100 / checks).append("%)"); + } else { + sb.append(" | OK"); + } + sb.append("\n"); + + long dur = s.totalDurationNanos.get(); + if (dur > 0 && checks > 0) { + long avgMs = (dur / checks) / 1_000_000; + long minMs = + s.minDurationNanos.get() == Long.MAX_VALUE ? 0 : s.minDurationNanos.get() / 1_000_000; + long maxMs = s.maxDurationNanos.get() / 1_000_000; + sb.append(" Latency: avg ").append(avgMs).append("ms"); + sb.append(", min ").append(minMs).append("ms"); + sb.append(", max ").append(maxMs).append("ms\n"); + } + + if (violations > 0 && s.lastViolationMessage != null) { + sb.append(" Last: ").append(s.lastViolationMessage).append("\n"); + } + sb.append("\n"); + } + + if (!taggedPaths.isEmpty()) { + sb.append("--- Tracked Code Paths ---\n"); + for (Map.Entry> e : taggedPaths.entrySet()) { + String contractName = e.getKey(); + Map byTag = e.getValue(); + + List tags = new ArrayList<>(byTag.keySet()); + Collections.sort(tags); + + sb.append(" ").append(contractName).append(":\n"); + for (String tag : tags) { + CodePathStats s = byTag.get(tag); + long calls = s.calls.get(); + long avgMs = calls > 0 ? (s.totalDurationNanos.get() / calls) / 1_000_000 : 0; + sb.append(" ") + .append(tag) + .append(": ") + .append(calls) + .append(" calls avg ") + .append(avgMs) + .append("ms\n"); + } + + // Cross-compare when exactly 2 tags + if (tags.size() == 2) { + CodePathStats s0 = byTag.get(tags.get(0)); + CodePathStats s1 = byTag.get(tags.get(1)); + long c0 = s0.calls.get(), c1 = s1.calls.get(); + if (c0 > 0 && c1 > 0) { + long avg0 = s0.totalDurationNanos.get() / c0; + long avg1 = s1.totalDurationNanos.get() / c1; + if (avg1 > 0) { + long pct = ((avg0 - avg1) * 100) / avg1; + if (pct > 0) { + sb.append(" [") + .append(tags.get(0)) + .append(' ') + .append(pct) + .append("% slower than ") + .append(tags.get(1)) + .append("]\n"); + } else if (pct < 0) { + sb.append(" [") + .append(tags.get(0)) + .append(' ') + .append(-pct) + .append("% faster than ") + .append(tags.get(1)) + .append("]\n"); + } + } + } + } + } + sb.append("\n"); + } + + sb.append("--- Totals ---\n"); + sb.append(" Checks: ").append(totalChecks).append("\n"); + sb.append(" Violations: ").append(totalViolations); + if (totalViolations == 0) { + sb.append(" (all contracts satisfied)"); + } + sb.append("\n"); + + return sb.toString(); + } + + @Override + public long getTotalViolations() { + long total = 0; + for (ContractStats s : contracts.values()) { + total += s.violations.get(); + } + return total; + } + + @Override + public long getViolations(String contract) { + ContractStats s = contracts.get(contract); + return s != null ? s.violations.get() : 0; + } + + @Override + public long getTotalChecks() { + long total = 0; + for (ContractStats s : contracts.values()) { + total += s.checks.get(); + } + return total; + } + + @Override + public boolean hasViolations() { + for (ContractStats s : contracts.values()) { + if (s.violations.get() > 0) return true; + } + return false; + } + + @Override + public void reset() { + contracts.clear(); + rateWindows.clear(); + taggedPaths.clear(); + } + + @Override + public void close() { + String summary = getSummary(); + if (!"No contracts checked.".equals(summary)) { + getContext().send(summary); + } + } + + // ==================== Internals ==================== + + private ContractStats getOrCreate(String contract) { + return contracts.computeIfAbsent(contract, k -> new ContractStats()); + } + + private static void updateMinMax(AtomicLong min, AtomicLong max, long value) { + long cur; + do { + cur = min.get(); + if (value >= cur) break; + } while (!min.compareAndSet(cur, value)); + do { + cur = max.get(); + if (value <= cur) break; + } while (!max.compareAndSet(cur, value)); + } + + static final class ContractStats { + final AtomicLong checks = new AtomicLong(); + final AtomicLong violations = new AtomicLong(); + final AtomicLong totalDurationNanos = new AtomicLong(); + final AtomicLong minDurationNanos = new AtomicLong(Long.MAX_VALUE); + final AtomicLong maxDurationNanos = new AtomicLong(0); + volatile String lastViolationMessage; + } + + static final class CodePathStats { + final AtomicLong calls = new AtomicLong(); + final AtomicLong totalDurationNanos = new AtomicLong(); + } + + /** + * Simple sliding-window rate counter. Tracks calls in the last second using a circular buffer of + * 10 x 100ms buckets. Lock-free via CAS. + */ + static final class RateWindow { + private static final int BUCKETS = 10; + private static final long BUCKET_NS = 100_000_000L; // 100ms + private final AtomicLong[] counts = new AtomicLong[BUCKETS]; + private final AtomicLong[] timestamps = new AtomicLong[BUCKETS]; + + RateWindow() { + for (int i = 0; i < BUCKETS; i++) { + counts[i] = new AtomicLong(); + timestamps[i] = new AtomicLong(); + } + } + + long recordAndCount(long nowNanos) { + int bucket = (int) ((nowNanos / BUCKET_NS) % BUCKETS); + long bucketTime = (nowNanos / BUCKET_NS) * BUCKET_NS; + + if (timestamps[bucket].get() != bucketTime) { + timestamps[bucket].set(bucketTime); + counts[bucket].set(0); + } + counts[bucket].incrementAndGet(); + + long total = 0; + long windowStart = nowNanos - (BUCKETS * BUCKET_NS); + for (int i = 0; i < BUCKETS; i++) { + if (timestamps[i].get() > windowStart) { + total += counts[i].get(); + } + } + return total; + } + } +} diff --git a/btrace-extensions/btrace-contracts/src/test/java/io/btrace/contracts/ContractServiceTest.java b/btrace-extensions/btrace-contracts/src/test/java/io/btrace/contracts/ContractServiceTest.java new file mode 100644 index 000000000..aa8a43488 --- /dev/null +++ b/btrace-extensions/btrace-contracts/src/test/java/io/btrace/contracts/ContractServiceTest.java @@ -0,0 +1,260 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.contracts; + +import static org.junit.jupiter.api.Assertions.*; + +import java.util.concurrent.CountDownLatch; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class ContractServiceTest { + + private ContractServiceImpl service; + + @BeforeEach + void setUp() { + service = new ContractServiceImpl(); + } + + // ==================== Latency checks ==================== + + @Test + void latencyWithinBudget() { + service.checkLatency("api.generate", 100_000_000L, 500_000_000L); + assertEquals(1, service.getTotalChecks()); + assertEquals(0, service.getTotalViolations()); + assertFalse(service.hasViolations()); + } + + @Test + void latencyExceedsBudget() { + service.checkLatency("api.generate", 600_000_000L, 500_000_000L); + assertEquals(1, service.getTotalViolations()); + assertTrue(service.hasViolations()); + String summary = service.getSummary(); + assertTrue(summary.contains("VIOLATIONS: 1")); + assertTrue(summary.contains("exceeded budget")); + } + + @Test + void latencyStatsTracked() { + service.checkLatency("api.call", 100_000_000L, 1_000_000_000L); + service.checkLatency("api.call", 200_000_000L, 1_000_000_000L); + service.checkLatency("api.call", 300_000_000L, 1_000_000_000L); + + assertEquals(3, service.getTotalChecks()); + String summary = service.getSummary(); + assertTrue(summary.contains("avg 200ms")); + assertTrue(summary.contains("min 100ms")); + assertTrue(summary.contains("max 300ms")); + } + + // ==================== Call rate checks ==================== + + @Test + void callRateWithinLimit() { + service.checkCallRate("api.query", 1000); + assertEquals(0, service.getTotalViolations()); + } + + @Test + void callRateExceedsLimit() { + for (int i = 0; i < 200; i++) { + service.checkCallRate("api.query", 10); + } + assertTrue(service.getTotalViolations() > 0); + } + + // ==================== Condition assertions ==================== + + @Test + void assertConditionTrue() { + service.assertCondition("positive-balance", true, "Balance must be positive"); + assertEquals(0, service.getTotalViolations()); + } + + @Test + void assertConditionFalse() { + service.assertCondition("positive-balance", false, "Balance must be positive"); + assertEquals(1, service.getTotalViolations()); + assertTrue(service.getSummary().contains("Balance must be positive")); + } + + // ==================== Range checks ==================== + + @Test + void rangeWithinBounds() { + service.checkRange("response-code", 200, 100, 599); + assertEquals(0, service.getTotalViolations()); + } + + @Test + void rangeBelowMin() { + service.checkRange("response-code", 50, 100, 599); + assertEquals(1, service.getTotalViolations()); + assertTrue(service.getSummary().contains("outside range")); + } + + @Test + void rangeAboveMax() { + service.checkRange("response-code", 700, 100, 599); + assertEquals(1, service.getTotalViolations()); + } + + @Test + void rangeAtBoundaries() { + service.checkRange("val", 100, 100, 200); + service.checkRange("val", 200, 100, 200); + assertEquals(0, service.getTotalViolations()); + } + + // ==================== Null checks ==================== + + @Test + void checkNotNullWithValue() { + service.checkNotNull("api.result", "hello"); + assertEquals(0, service.getTotalViolations()); + } + + @Test + void checkNotNullWithNull() { + service.checkNotNull("api.result", null); + assertEquals(1, service.getTotalViolations()); + assertTrue(service.getSummary().contains("Unexpected null")); + } + + // ==================== Code path tracking ==================== + + @Test + void twoTagComparison() { + service.trackCodePath("Parser.parse", 50_000_000L, "v2"); + service.trackCodePath("Parser.parse", 60_000_000L, "v2"); + service.trackCodePath("Parser.parse", 30_000_000L, "v1"); + service.trackCodePath("Parser.parse", 40_000_000L, "v1"); + + String summary = service.getSummary(); + assertTrue(summary.contains("Tracked Code Paths")); + assertTrue(summary.contains("v1")); + assertTrue(summary.contains("v2")); + assertTrue(summary.contains("2 calls")); + // v2 avg 55ms, v1 avg 35ms -> v2 ~57% slower than v1 (alphabetical: v1 first, v2 second) + assertTrue(summary.contains("slower") || summary.contains("faster")); + } + + @Test + void singleTagTracking() { + service.trackCodePath("Renderer.render", 100_000_000L, "cached"); + String summary = service.getSummary(); + assertTrue(summary.contains("Tracked Code Paths")); + assertTrue(summary.contains("cached")); + assertTrue(summary.contains("1 calls")); + } + + @Test + void multipleTagsNoComparison() { + service.trackCodePath("Op.run", 10_000_000L, "a"); + service.trackCodePath("Op.run", 20_000_000L, "b"); + service.trackCodePath("Op.run", 30_000_000L, "c"); + + String summary = service.getSummary(); + assertTrue(summary.contains("a")); + assertTrue(summary.contains("b")); + assertTrue(summary.contains("c")); + // More than 2 tags: no slower/faster comparison + assertFalse( + summary.contains("slower") || summary.contains("faster"), + "Cross-comparison should only appear for exactly 2 tags"); + } + + // ==================== Reporting ==================== + + @Test + void noDataSummary() { + assertEquals("No contracts checked.", service.getSummary()); + } + + @Test + void getViolationsPerContract() { + service.assertCondition("a", false, "fail"); + service.assertCondition("a", false, "fail"); + service.assertCondition("b", false, "fail"); + + assertEquals(2, service.getViolations("a")); + assertEquals(1, service.getViolations("b")); + assertEquals(0, service.getViolations("c")); + } + + @Test + void multipleContractsInSummary() { + service.checkLatency("fast-api", 10_000_000L, 100_000_000L); + service.checkLatency("slow-api", 500_000_000L, 100_000_000L); + + String summary = service.getSummary(); + assertTrue(summary.contains("fast-api")); + assertTrue(summary.contains("slow-api")); + assertTrue(summary.contains("Checks: 2")); + assertTrue(summary.contains("Violations: 1")); + } + + @Test + void allContractsSatisfied() { + service.checkLatency("api", 10_000_000L, 100_000_000L); + service.assertCondition("invariant", true, "ok"); + service.checkRange("val", 50, 0, 100); + + String summary = service.getSummary(); + assertTrue(summary.contains("all contracts satisfied")); + } + + @Test + void reset() { + service.assertCondition("a", false, "fail"); + service.trackCodePath("b", 100L, "tag"); + service.reset(); + + assertEquals(0, service.getTotalChecks()); + assertEquals(0, service.getTotalViolations()); + assertFalse(service.hasViolations()); + assertEquals("No contracts checked.", service.getSummary()); + } + + @Test + void concurrentChecks() throws Exception { + int threads = 8; + int checksPerThread = 1000; + CountDownLatch latch = new CountDownLatch(threads); + + for (int t = 0; t < threads; t++) { + new Thread( + () -> { + try { + for (int i = 0; i < checksPerThread; i++) { + service.checkLatency("concurrent-api", 50_000_000L, 100_000_000L); + } + } finally { + latch.countDown(); + } + }) + .start(); + } + latch.await(); + + assertEquals(threads * checksPerThread, service.getTotalChecks()); + assertEquals(0, service.getTotalViolations()); + } +} diff --git a/btrace-extensions/btrace-gpu-bridge/build.gradle b/btrace-extensions/btrace-gpu-bridge/build.gradle new file mode 100644 index 000000000..254087a79 --- /dev/null +++ b/btrace-extensions/btrace-gpu-bridge/build.gradle @@ -0,0 +1,35 @@ +plugins { + id 'io.btrace.extension' + alias(libs.plugins.shadow) +} + +java { + sourceCompatibility = 8 + targetCompatibility = 8 +} + +compileJava { + javaCompiler = javaToolchains.compilerFor { + languageVersion.set(JavaLanguageVersion.of(11)) + } +} + +btraceExtension { + id = 'btrace-gpu-bridge' + name = 'BTrace GPU Bridge' + description = 'GPU and model inference observability: ONNX Runtime, DJL, TensorFlow Java, and Panama FFM compute tracing' + services = ['io.btrace.gpu.GpuBridgeService'] +} + +dependencies { + apiCompileOnly 'org.jetbrains:annotations:26.1.0' + apiCompileOnly project(':btrace-core') + implCompileOnly 'org.jetbrains:annotations:26.1.0' + implImplementation project(':btrace-core') + + testImplementation libs.junit.jupiter +} + +test { + useJUnitPlatform() +} diff --git a/btrace-extensions/btrace-gpu-bridge/src/main/java/io/btrace/gpu/GpuBridgeService.java b/btrace-extensions/btrace-gpu-bridge/src/main/java/io/btrace/gpu/GpuBridgeService.java new file mode 100644 index 000000000..44c101a55 --- /dev/null +++ b/btrace-extensions/btrace-gpu-bridge/src/main/java/io/btrace/gpu/GpuBridgeService.java @@ -0,0 +1,148 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.gpu; + +import io.btrace.core.extensions.ServiceDescriptor; +import org.jetbrains.annotations.Nullable; + +/** + * BTrace extension for GPU compute and model inference observability. + * + *

Traces Java-to-GPU boundaries: ONNX Runtime inference sessions, DJL (Deep Java Library) + * predictions, TensorFlow Java operations, and Panama FFM native calls to CUDA/ROCm libraries. Zero + * external dependencies — instruments existing client classes. + * + *

Usage in a BTrace script: + * + * 

+ * @Injected GpuBridgeService gpu;
+ *
+ * @OnMethod(clazz = "ai.onnxruntime.OrtSession", method = "run")
+ * void onInference(@Duration long dur) {
+ *     gpu.recordInference("onnx", "resnet50", dur);
+ * }
+ *
+ */ +@ServiceDescriptor +public interface GpuBridgeService { + + // ==================== Inference recording ==================== + + /** + * Records a model inference call with duration only. + * + * @param runtime runtime name (e.g. "onnx", "djl", "tensorflow") + * @param modelName model identifier + * @param durationNanos inference duration in nanoseconds + */ + void recordInference(String runtime, String modelName, long durationNanos); + + /** + * Records a model inference with batch size and tensor dimensions. + * + * @param runtime runtime name + * @param modelName model identifier + * @param batchSize batch size of the inference request + * @param durationNanos inference duration + */ + void recordInference(String runtime, String modelName, int batchSize, long durationNanos); + + /** + * Starts a detailed inference record builder. Allocation-free (ThreadLocal-pooled). + * + * 
+   * gpu.inference("onnx", "bert-base")
+   *     .batchSize(32)
+   *     .inputElements(512 * 768)
+   *     .outputElements(512 * 2)
+   *     .deviceType("cuda")
+   *     .deviceId(0)
+   *     .duration(durationNanos)
+   *     .record();
+   *
+ * + * @param runtime runtime name + * @param modelName model identifier + * @return an inference record builder (thread-local, do not store) + */ + @Nullable InferenceRecord inference(String runtime, String modelName); + + // ==================== Memory tracking ==================== + + /** + * Records GPU memory allocation. + * + * @param deviceType device type ("cuda", "rocm", "cpu") + * @param deviceId device index + * @param bytes allocated bytes + */ + void recordMemoryAlloc(String deviceType, int deviceId, long bytes); + + /** + * Records GPU memory deallocation. + * + * @param deviceType device type + * @param deviceId device index + * @param bytes freed bytes + */ + void recordMemoryFree(String deviceType, int deviceId, long bytes); + + // ==================== Native call tracking ==================== + + /** + * Records a native/FFM call to a GPU library (e.g. cuBLAS, cuDNN). + * + * @param library library name + * @param function function name + * @param durationNanos call duration + */ + void recordNativeCall(String library, String function, long durationNanos); + + // ==================== Model lifecycle ==================== + + /** + * Records model load/initialization time. + * + * @param runtime runtime name + * @param modelName model identifier + * @param durationNanos load duration + */ + void recordModelLoad(String runtime, String modelName, long durationNanos); + + // ==================== Reporting ==================== + + /** Returns a formatted summary of all GPU/inference metrics. */ + @Nullable String getSummary(); + + /** Returns summary for a specific model. */ + @Nullable String getModelSummary(String modelName); + + /** Total number of inference calls across all models. */ + long getTotalInferences(); + + /** Returns estimated GPU memory currently allocated (bytes), or -1 if not tracked. */ + long getCurrentGpuMemoryBytes(); + + /** Returns peak GPU memory seen (bytes), or -1 if not tracked. */ + long getPeakGpuMemoryBytes(); + + /** Returns total number of native/FFM calls recorded. */ + long getTotalNativeCalls(); + + /** Resets all collected metrics. */ + void reset(); +} diff --git a/btrace-extensions/btrace-gpu-bridge/src/main/java/io/btrace/gpu/GpuBridgeServiceImpl.java b/btrace-extensions/btrace-gpu-bridge/src/main/java/io/btrace/gpu/GpuBridgeServiceImpl.java new file mode 100644 index 000000000..1eaee7337 --- /dev/null +++ b/btrace-extensions/btrace-gpu-bridge/src/main/java/io/btrace/gpu/GpuBridgeServiceImpl.java @@ -0,0 +1,360 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.gpu; + +import io.btrace.core.extensions.Extension; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicLong; + +/** Thread-safe GPU compute and inference tracking with lock-free statistics. */ +public final class GpuBridgeServiceImpl extends Extension implements GpuBridgeService { + + private final Map modelStats = new ConcurrentHashMap<>(); + private final Map deviceMemory = new ConcurrentHashMap<>(); + private final Map nativeStats = new ConcurrentHashMap<>(); + + private final ThreadLocal inferenceRecordPool = + ThreadLocal.withInitial(InferenceRecordImpl::new); + + // ==================== Inference recording ==================== + + @Override + public void recordInference(String runtime, String modelName, long durationNanos) { + ModelStats stats = getOrCreateModel(runtime, modelName); + stats.inferences.incrementAndGet(); + stats.totalDurationNanos.addAndGet(durationNanos); + updateMinMax(stats.minDurationNanos, stats.maxDurationNanos, durationNanos); + } + + @Override + public void recordInference(String runtime, String modelName, int batchSize, long durationNanos) { + ModelStats stats = getOrCreateModel(runtime, modelName); + stats.inferences.incrementAndGet(); + stats.totalBatchSize.addAndGet(batchSize); + stats.totalDurationNanos.addAndGet(durationNanos); + updateMinMax(stats.minDurationNanos, stats.maxDurationNanos, durationNanos); + } + + @Override + public InferenceRecord inference(String runtime, String modelName) { + return inferenceRecordPool.get().reset(this, runtime, modelName); + } + + void commitInferenceRecord(InferenceRecordImpl rec) { + ModelStats stats = getOrCreateModel(rec.runtime, rec.modelName); + stats.inferences.incrementAndGet(); + stats.totalBatchSize.addAndGet(rec.batchSizeVal); + stats.totalInputElements.addAndGet(rec.inputElem); + stats.totalOutputElements.addAndGet(rec.outputElem); + stats.totalDurationNanos.addAndGet(rec.durationVal); + updateMinMax(stats.minDurationNanos, stats.maxDurationNanos, rec.durationVal); + if (rec.deviceTypeVal != null) { + stats.lastDeviceType = rec.deviceTypeVal; + stats.lastDeviceId = rec.deviceIdVal; + } + } + + // ==================== Memory tracking ==================== + + @Override + public void recordMemoryAlloc(String deviceType, int deviceId, long bytes) { + DeviceMemory dm = getOrCreateDevice(deviceType, deviceId); + dm.currentBytes.addAndGet(bytes); + dm.totalAllocated.addAndGet(bytes); + dm.allocCount.incrementAndGet(); + // Update peak + long cur; + long newVal = dm.currentBytes.get(); + do { + cur = dm.peakBytes.get(); + if (newVal <= cur) break; + } while (!dm.peakBytes.compareAndSet(cur, newVal)); + } + + @Override + public void recordMemoryFree(String deviceType, int deviceId, long bytes) { + DeviceMemory dm = getOrCreateDevice(deviceType, deviceId); + dm.currentBytes.addAndGet(-bytes); + dm.freeCount.incrementAndGet(); + } + + // ==================== Native call tracking ==================== + + @Override + public void recordNativeCall(String library, String function, long durationNanos) { + String key = library + "::" + function; + NativeCallStats stats = nativeStats.computeIfAbsent(key, k -> new NativeCallStats()); + stats.calls.incrementAndGet(); + stats.totalDurationNanos.addAndGet(durationNanos); + } + + // ==================== Model lifecycle ==================== + + @Override + public void recordModelLoad(String runtime, String modelName, long durationNanos) { + ModelStats stats = getOrCreateModel(runtime, modelName); + stats.loadCount.incrementAndGet(); + stats.totalLoadDurationNanos.addAndGet(durationNanos); + } + + // ==================== Reporting ==================== + + @Override + public String getSummary() { + if (modelStats.isEmpty() && deviceMemory.isEmpty() && nativeStats.isEmpty()) { + return "No GPU/inference activity recorded."; + } + + StringBuilder sb = new StringBuilder(); + sb.append("=== GPU Bridge Summary ===\n\n"); + + long totalInf = 0; + + for (Map.Entry entry : modelStats.entrySet()) { + ModelStats s = entry.getValue(); + long inf = s.inferences.get(); + totalInf += inf; + + sb.append("Model: ").append(entry.getKey()); + if (s.lastDeviceType != null) { + sb.append(" [").append(s.lastDeviceType).append(":").append(s.lastDeviceId).append("]"); + } + sb.append("\n"); + + // Inference stats + sb.append(" Inferences: ").append(inf); + long loads = s.loadCount.get(); + if (loads > 0) { + long avgLoadMs = (s.totalLoadDurationNanos.get() / loads) / 1_000_000; + sb.append(" (loaded ").append(loads).append("x, avg ").append(avgLoadMs).append("ms)"); + } + sb.append("\n"); + + // Latency + if (inf > 0) { + long avgMs = (s.totalDurationNanos.get() / inf) / 1_000_000; + long minMs = s.minDurationNanos.get() / 1_000_000; + long maxMs = s.maxDurationNanos.get() / 1_000_000; + sb.append(" Latency: avg ").append(avgMs).append("ms"); + sb.append(", min ").append(minMs).append("ms"); + sb.append(", max ").append(maxMs).append("ms\n"); + } + + // Batch size + long totalBatch = s.totalBatchSize.get(); + if (totalBatch > 0 && inf > 0) { + sb.append(" Batch size: avg ").append(totalBatch / inf).append("\n"); + } + + // Throughput (items/sec) + long totalDur = s.totalDurationNanos.get(); + if (totalBatch > 0 && totalDur > 0) { + double itemsPerSec = (totalBatch * 1_000_000_000.0) / totalDur; + sb.append(" Throughput: ") + .append(String.format("%.1f", itemsPerSec)) + .append(" items/sec\n"); + } + + // Tensor elements + long inEl = s.totalInputElements.get(); + long outEl = s.totalOutputElements.get(); + if (inEl > 0 || outEl > 0) { + sb.append(" Tensors: ") + .append(formatElements(inEl)) + .append(" in / ") + .append(formatElements(outEl)) + .append(" out\n"); + } + + sb.append("\n"); + } + + // Device memory + if (!deviceMemory.isEmpty()) { + sb.append("--- Device Memory ---\n"); + for (Map.Entry entry : deviceMemory.entrySet()) { + DeviceMemory dm = entry.getValue(); + sb.append(" ").append(entry.getKey()).append(": "); + sb.append("current ").append(formatBytes(dm.currentBytes.get())); + sb.append(", peak ").append(formatBytes(dm.peakBytes.get())); + sb.append(" (") + .append(dm.allocCount.get()) + .append(" allocs, ") + .append(dm.freeCount.get()) + .append(" frees)\n"); + } + sb.append("\n"); + } + + // Native calls + if (!nativeStats.isEmpty()) { + sb.append("--- Native Calls ---\n"); + for (Map.Entry entry : nativeStats.entrySet()) { + NativeCallStats ns = entry.getValue(); + long calls = ns.calls.get(); + long avgUs = calls > 0 ? (ns.totalDurationNanos.get() / calls) / 1000 : 0; + sb.append(" ") + .append(entry.getKey()) + .append(": ") + .append(calls) + .append(" calls, avg ") + .append(avgUs) + .append("us\n"); + } + sb.append("\n"); + } + + sb.append("--- Totals ---\n"); + sb.append(" Inferences: ").append(totalInf).append("\n"); + + return sb.toString(); + } + + @Override + public String getModelSummary(String modelName) { + // Search by model name suffix + for (Map.Entry entry : modelStats.entrySet()) { + if (entry.getKey().endsWith(modelName) || entry.getKey().equals(modelName)) { + ModelStats s = entry.getValue(); + long inf = s.inferences.get(); + long avgMs = inf > 0 ? (s.totalDurationNanos.get() / inf) / 1_000_000 : 0; + return entry.getKey() + ": " + inf + " inferences, avg " + avgMs + "ms"; + } + } + return "No data for model: " + modelName; + } + + @Override + public long getTotalInferences() { + long total = 0; + for (ModelStats s : modelStats.values()) { + total += s.inferences.get(); + } + return total; + } + + @Override + public long getCurrentGpuMemoryBytes() { + if (deviceMemory.isEmpty()) return -1; + long total = 0; + for (DeviceMemory dm : deviceMemory.values()) { + total += dm.currentBytes.get(); + } + return total; + } + + @Override + public long getPeakGpuMemoryBytes() { + if (deviceMemory.isEmpty()) return -1; + long peak = 0; + for (DeviceMemory dm : deviceMemory.values()) { + peak = Math.max(peak, dm.peakBytes.get()); + } + return peak; + } + + @Override + public long getTotalNativeCalls() { + long total = 0; + for (NativeCallStats ns : nativeStats.values()) { + total += ns.calls.get(); + } + return total; + } + + @Override + public void reset() { + modelStats.clear(); + deviceMemory.clear(); + nativeStats.clear(); + } + + @Override + public void close() { + String summary = getSummary(); + if (!"No GPU/inference activity recorded.".equals(summary)) { + getContext().send(summary); + } + } + + // ==================== Internals ==================== + + private ModelStats getOrCreateModel(String runtime, String modelName) { + String key = runtime + "/" + modelName; + return modelStats.computeIfAbsent(key, k -> new ModelStats()); + } + + private DeviceMemory getOrCreateDevice(String deviceType, int deviceId) { + String key = deviceType + ":" + deviceId; + return deviceMemory.computeIfAbsent(key, k -> new DeviceMemory()); + } + + private static void updateMinMax(AtomicLong min, AtomicLong max, long value) { + long cur; + do { + cur = min.get(); + if (value >= cur) break; + } while (!min.compareAndSet(cur, value)); + do { + cur = max.get(); + if (value <= cur) break; + } while (!max.compareAndSet(cur, value)); + } + + static String formatBytes(long bytes) { + if (bytes < 0) return "-" + formatBytes(-bytes); + if (bytes < 1024) return bytes + "B"; + if (bytes < 1024 * 1024) return (bytes / 1024) + "KB"; + if (bytes < 1024L * 1024 * 1024) return (bytes / (1024 * 1024)) + "MB"; + return String.format("%.1fGB", bytes / (1024.0 * 1024 * 1024)); + } + + static String formatElements(long elements) { + if (elements < 1000) return elements + ""; + if (elements < 1_000_000) return String.format("%.1fK", elements / 1000.0); + if (elements < 1_000_000_000) return String.format("%.1fM", elements / 1_000_000.0); + return String.format("%.1fB", elements / 1_000_000_000.0); + } + + static final class ModelStats { + final AtomicLong inferences = new AtomicLong(); + final AtomicLong totalBatchSize = new AtomicLong(); + final AtomicLong totalInputElements = new AtomicLong(); + final AtomicLong totalOutputElements = new AtomicLong(); + final AtomicLong totalDurationNanos = new AtomicLong(); + final AtomicLong minDurationNanos = new AtomicLong(Long.MAX_VALUE); + final AtomicLong maxDurationNanos = new AtomicLong(0); + final AtomicLong loadCount = new AtomicLong(); + final AtomicLong totalLoadDurationNanos = new AtomicLong(); + volatile String lastDeviceType; + volatile int lastDeviceId; + } + + static final class DeviceMemory { + final AtomicLong currentBytes = new AtomicLong(); + final AtomicLong peakBytes = new AtomicLong(); + final AtomicLong totalAllocated = new AtomicLong(); + final AtomicLong allocCount = new AtomicLong(); + final AtomicLong freeCount = new AtomicLong(); + } + + static final class NativeCallStats { + final AtomicLong calls = new AtomicLong(); + final AtomicLong totalDurationNanos = new AtomicLong(); + } +} diff --git a/btrace-extensions/btrace-gpu-bridge/src/main/java/io/btrace/gpu/InferenceRecord.java b/btrace-extensions/btrace-gpu-bridge/src/main/java/io/btrace/gpu/InferenceRecord.java new file mode 100644 index 000000000..ca848f35d --- /dev/null +++ b/btrace-extensions/btrace-gpu-bridge/src/main/java/io/btrace/gpu/InferenceRecord.java @@ -0,0 +1,58 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.gpu; + +/** + * Fluent builder for recording a model inference with detailed metrics. + * + *

Obtain via {@link GpuBridgeService#inference(String, String)}. Allocation-free + * (ThreadLocal-pooled). Do not store the returned reference. + * + * 

+ * gpu.inference("onnx", "bert-base")
+ *     .batchSize(32)
+ *     .inputElements(512 * 768)
+ *     .outputElements(512 * 2)
+ *     .deviceType("cuda")
+ *     .deviceId(0)
+ *     .duration(durationNanos)
+ *     .record();
+ *
+ */ +public interface InferenceRecord { + + /** Batch size for this inference. */ + InferenceRecord batchSize(int size); + + /** Total number of input tensor elements. */ + InferenceRecord inputElements(long elements); + + /** Total number of output tensor elements. */ + InferenceRecord outputElements(long elements); + + /** Device type: "cuda", "rocm", "cpu", "mps", etc. */ + InferenceRecord deviceType(String type); + + /** Device index (for multi-GPU). */ + InferenceRecord deviceId(int id); + + /** Inference duration in nanoseconds. */ + InferenceRecord duration(long nanos); + + /** Commits this inference record. */ + void record(); +} diff --git a/btrace-extensions/btrace-gpu-bridge/src/main/java/io/btrace/gpu/InferenceRecordImpl.java b/btrace-extensions/btrace-gpu-bridge/src/main/java/io/btrace/gpu/InferenceRecordImpl.java new file mode 100644 index 000000000..834621cac --- /dev/null +++ b/btrace-extensions/btrace-gpu-bridge/src/main/java/io/btrace/gpu/InferenceRecordImpl.java @@ -0,0 +1,88 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.gpu; + +/** ThreadLocal-pooled builder for {@link InferenceRecord}. Zero allocation. */ +final class InferenceRecordImpl implements InferenceRecord { + + String runtime; + String modelName; + int batchSizeVal; + long inputElem; + long outputElem; + String deviceTypeVal; + int deviceIdVal; + long durationVal; + + private GpuBridgeServiceImpl service; + + InferenceRecordImpl() {} + + InferenceRecordImpl reset(GpuBridgeServiceImpl service, String runtime, String modelName) { + this.service = service; + this.runtime = runtime; + this.modelName = modelName; + this.batchSizeVal = 1; + this.inputElem = 0; + this.outputElem = 0; + this.deviceTypeVal = null; + this.deviceIdVal = 0; + this.durationVal = 0; + return this; + } + + @Override + public InferenceRecord batchSize(int size) { + this.batchSizeVal = size; + return this; + } + + @Override + public InferenceRecord inputElements(long elements) { + this.inputElem = elements; + return this; + } + + @Override + public InferenceRecord outputElements(long elements) { + this.outputElem = elements; + return this; + } + + @Override + public InferenceRecord deviceType(String type) { + this.deviceTypeVal = type; + return this; + } + + @Override + public InferenceRecord deviceId(int id) { + this.deviceIdVal = id; + return this; + } + + @Override + public InferenceRecord duration(long nanos) { + this.durationVal = nanos; + return this; + } + + @Override + public void record() { + service.commitInferenceRecord(this); + } +} diff --git a/btrace-extensions/btrace-gpu-bridge/src/test/java/io/btrace/gpu/GpuBridgeServiceTest.java b/btrace-extensions/btrace-gpu-bridge/src/test/java/io/btrace/gpu/GpuBridgeServiceTest.java new file mode 100644 index 000000000..63f493115 --- /dev/null +++ b/btrace-extensions/btrace-gpu-bridge/src/test/java/io/btrace/gpu/GpuBridgeServiceTest.java @@ -0,0 +1,238 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.gpu; + +import static org.junit.jupiter.api.Assertions.*; + +import java.util.concurrent.CountDownLatch; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class GpuBridgeServiceTest { + + private GpuBridgeServiceImpl service; + + @BeforeEach + void setUp() { + service = new GpuBridgeServiceImpl(); + } + + // ==================== Inference recording ==================== + + @Test + void durationOnlyInference() { + service.recordInference("onnx", "resnet50", 10_000_000L); + assertEquals(1, service.getTotalInferences()); + assertTrue(service.getSummary().contains("onnx/resnet50")); + } + + @Test + void inferenceWithBatchSize() { + service.recordInference("djl", "bert", 32, 50_000_000L); + assertEquals(1, service.getTotalInferences()); + assertTrue(service.getSummary().contains("avg 32")); + } + + @Test + void fluentBuilder() { + service + .inference("onnx", "yolo-v8") + .batchSize(16) + .inputElements(640L * 640 * 3) + .outputElements(8400L * 84) + .deviceType("cuda") + .deviceId(0) + .duration(25_000_000L) + .record(); + + assertEquals(1, service.getTotalInferences()); + String summary = service.getSummary(); + assertTrue(summary.contains("cuda:0")); + assertTrue(summary.contains("items/sec")); + } + + @Test + void fluentBuilderMinimal() { + service.inference("tensorflow", "mobilenet").duration(5_000_000L).record(); + + assertEquals(1, service.getTotalInferences()); + } + + @Test + void multipleModelsTrackedSeparately() { + service.recordInference("onnx", "resnet50", 10_000_000L); + service.recordInference("onnx", "bert", 20_000_000L); + service.recordInference("onnx", "resnet50", 12_000_000L); + + assertEquals(3, service.getTotalInferences()); + assertTrue(service.getModelSummary("resnet50").contains("2 inferences")); + assertTrue(service.getModelSummary("bert").contains("1 inferences")); + } + + @Test + void latencyMinMax() { + service.recordInference("onnx", "model", 5_000_000L); + service.recordInference("onnx", "model", 50_000_000L); + service.recordInference("onnx", "model", 10_000_000L); + + String summary = service.getSummary(); + assertTrue(summary.contains("min 5ms")); + assertTrue(summary.contains("max 50ms")); + } + + // ==================== Memory tracking ==================== + + @Test + void memoryAllocAndFree() { + service.recordMemoryAlloc("cuda", 0, 1024 * 1024 * 100); // 100MB + assertEquals(100 * 1024 * 1024L, service.getCurrentGpuMemoryBytes()); + assertEquals(100 * 1024 * 1024L, service.getPeakGpuMemoryBytes()); + + service.recordMemoryFree("cuda", 0, 1024 * 1024 * 50); // free 50MB + assertEquals(50 * 1024 * 1024L, service.getCurrentGpuMemoryBytes()); + assertEquals(100 * 1024 * 1024L, service.getPeakGpuMemoryBytes()); // peak unchanged + } + + @Test + void noMemoryTracked() { + assertEquals(-1, service.getCurrentGpuMemoryBytes()); + assertEquals(-1, service.getPeakGpuMemoryBytes()); + } + + @Test + void memorySummary() { + service.recordMemoryAlloc("cuda", 0, 500 * 1024 * 1024L); + String summary = service.getSummary(); + assertTrue(summary.contains("Device Memory")); + assertTrue(summary.contains("500MB")); + } + + // ==================== Native call tracking ==================== + + @Test + void nativeCalls() { + service.recordNativeCall("cublas", "sgemm", 500_000L); + service.recordNativeCall("cublas", "sgemm", 600_000L); + service.recordNativeCall("cudnn", "conv_forward", 1_000_000L); + + assertEquals(3, service.getTotalNativeCalls()); + String summary = service.getSummary(); + assertTrue(summary.contains("cublas::sgemm")); + assertTrue(summary.contains("2 calls")); + } + + // ==================== Model lifecycle ==================== + + @Test + void modelLoad() { + service.recordModelLoad("onnx", "bert", 2_000_000_000L); // 2 seconds + service.recordInference("onnx", "bert", 10_000_000L); + + String summary = service.getSummary(); + assertTrue(summary.contains("loaded 1x")); + assertTrue(summary.contains("2000ms")); + } + + // ==================== Reporting ==================== + + @Test + void noDataSummary() { + assertEquals("No GPU/inference activity recorded.", service.getSummary()); + } + + @Test + void unknownModelSummary() { + assertEquals("No data for model: unknown", service.getModelSummary("unknown")); + } + + @Test + void formatBytes() { + assertEquals("100B", GpuBridgeServiceImpl.formatBytes(100)); + assertEquals("10KB", GpuBridgeServiceImpl.formatBytes(10 * 1024)); + assertEquals("256MB", GpuBridgeServiceImpl.formatBytes(256 * 1024 * 1024)); + assertEquals("1.5GB", GpuBridgeServiceImpl.formatBytes((long) (1.5 * 1024 * 1024 * 1024))); + } + + @Test + void formatElements() { + assertEquals("500", GpuBridgeServiceImpl.formatElements(500)); + assertEquals("1.5K", GpuBridgeServiceImpl.formatElements(1500)); + assertEquals("2.0M", GpuBridgeServiceImpl.formatElements(2_000_000)); + assertEquals("1.0B", GpuBridgeServiceImpl.formatElements(1_000_000_000)); + } + + @Test + void reset() { + service.recordInference("onnx", "model", 1L); + service.recordMemoryAlloc("cuda", 0, 1024); + service.recordNativeCall("cublas", "fn", 1L); + service.reset(); + + assertEquals(0, service.getTotalInferences()); + assertEquals(-1, service.getCurrentGpuMemoryBytes()); + assertEquals(0, service.getTotalNativeCalls()); + assertEquals("No GPU/inference activity recorded.", service.getSummary()); + } + + @Test + void concurrentInference() throws Exception { + int threads = 8; + int infsPerThread = 1000; + CountDownLatch latch = new CountDownLatch(threads); + + for (int t = 0; t < threads; t++) { + new Thread( + () -> { + try { + for (int i = 0; i < infsPerThread; i++) { + service.recordInference("onnx", "model", 32, 1_000_000L); + } + } finally { + latch.countDown(); + } + }) + .start(); + } + latch.await(); + + assertEquals(threads * infsPerThread, service.getTotalInferences()); + } + + @Test + void concurrentBuilderInference() throws Exception { + int threads = 8; + int infsPerThread = 500; + CountDownLatch latch = new CountDownLatch(threads); + + for (int t = 0; t < threads; t++) { + new Thread( + () -> { + try { + for (int i = 0; i < infsPerThread; i++) { + service.inference("djl", "bert").batchSize(16).duration(2_000_000L).record(); + } + } finally { + latch.countDown(); + } + }) + .start(); + } + latch.await(); + + assertEquals(threads * infsPerThread, service.getTotalInferences()); + } +} diff --git a/btrace-extensions/btrace-llm-trace/build.gradle b/btrace-extensions/btrace-llm-trace/build.gradle new file mode 100644 index 000000000..a761069ff --- /dev/null +++ b/btrace-extensions/btrace-llm-trace/build.gradle @@ -0,0 +1,35 @@ +plugins { + id 'io.btrace.extension' + alias(libs.plugins.shadow) +} + +java { + sourceCompatibility = 8 + targetCompatibility = 8 +} + +compileJava { + javaCompiler = javaToolchains.compilerFor { + languageVersion.set(JavaLanguageVersion.of(11)) + } +} + +btraceExtension { + id = 'btrace-llm-trace' + name = 'BTrace LLM Tracing' + description = 'LLM inference observability: token counts, latency, cost tracking for Java LLM SDKs' + services = ['io.btrace.llm.LlmTraceService'] +} + +dependencies { + apiCompileOnly 'org.jetbrains:annotations:26.1.0' + apiCompileOnly project(':btrace-core') + implCompileOnly 'org.jetbrains:annotations:26.1.0' + implImplementation project(':btrace-core') + + testImplementation libs.junit.jupiter +} + +test { + useJUnitPlatform() +} diff --git a/btrace-extensions/btrace-llm-trace/src/main/java/io/btrace/llm/CallRecord.java b/btrace-extensions/btrace-llm-trace/src/main/java/io/btrace/llm/CallRecord.java new file mode 100644 index 000000000..07089e7be --- /dev/null +++ b/btrace-extensions/btrace-llm-trace/src/main/java/io/btrace/llm/CallRecord.java @@ -0,0 +1,88 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.llm; + +/** + * Fluent builder for recording an LLM API call. + * + *

Obtain via {@link LlmTraceService#call(String)}. All setters are optional except {@link + * #duration(long)} — if omitted, duration defaults to 0. + * + *

Call {@link #record()} to commit the metrics. + * + *

Allocation-free: Instances are pooled per-thread internally, so calling + * {@code call()} does not allocate on the heap. Safe for hot-path instrumentation. However, the + * returned reference must not be stored or shared across threads — use it inline and call {@link + * #record()} immediately. + * + * 

+ * llm.call("claude-sonnet-4-20250514")
+ *     .provider("anthropic")
+ *     .inputTokens(1500)
+ *     .outputTokens(300)
+ *     .cacheReadTokens(800)
+ *     .duration(durationNanos)
+ *     .record();
+ *
+ */ +public interface CallRecord { + + /** Sets the provider name (e.g. "openai", "anthropic", "google"). */ + CallRecord provider(String provider); + + /** Sets the number of input/prompt tokens. */ + CallRecord inputTokens(int tokens); + + /** Sets the number of output/completion tokens. */ + CallRecord outputTokens(int tokens); + + /** + * Sets the number of cache-read input tokens (prompt caching). Anthropic: {@code + * usage.cache_read_input_tokens}. OpenAI: {@code usage.prompt_tokens_details.cached_tokens}. + */ + CallRecord cacheReadTokens(int tokens); + + /** + * Sets the number of cache-creation input tokens. Anthropic: {@code + * usage.cache_creation_input_tokens}. + */ + CallRecord cacheCreationTokens(int tokens); + + /** Marks this call as a streaming response. */ + CallRecord streaming(); + + /** + * Sets the time-to-first-token for streaming calls. Only meaningful when {@link #streaming()} is + * set. + * + * @param nanos time from request start to first token, in nanoseconds + */ + CallRecord timeToFirstToken(long nanos); + + /** + * Sets the total call duration. + * + * @param nanos duration in nanoseconds (typically from {@code @Duration}) + */ + CallRecord duration(long nanos); + + /** + * Commits this call record to the trace service. Must be called exactly once to record the + * metrics. + */ + void record(); +} diff --git a/btrace-extensions/btrace-llm-trace/src/main/java/io/btrace/llm/CallRecordImpl.java b/btrace-extensions/btrace-llm-trace/src/main/java/io/btrace/llm/CallRecordImpl.java new file mode 100644 index 000000000..8906b8b0c --- /dev/null +++ b/btrace-extensions/btrace-llm-trace/src/main/java/io/btrace/llm/CallRecordImpl.java @@ -0,0 +1,114 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.llm; + +/** + * Mutable builder implementing the {@link CallRecord} fluent API. + * + *

Instances are pooled per-thread via {@link ThreadLocal} to avoid heap allocation on every + * {@code call()} invocation. This makes the builder safe for hot-path instrumentation — no garbage + * is produced. + * + *

Not thread-safe — intended to be created, configured, and recorded within a single BTrace + * handler method invocation on the same thread. + */ +final class CallRecordImpl implements CallRecord { + + String model; + String providerVal; + int inputTok; + int outputTok; + int cacheReadTok; + int cacheCreateTok; + boolean isStreaming; + long ttftVal; + long durationVal; + + private LlmTraceServiceImpl service; + + CallRecordImpl() { + // Created once per thread via ThreadLocal + } + + /** Resets all fields and binds this record to a new call. */ + CallRecordImpl reset(LlmTraceServiceImpl service, String model) { + this.service = service; + this.model = model; + this.providerVal = null; + this.inputTok = 0; + this.outputTok = 0; + this.cacheReadTok = 0; + this.cacheCreateTok = 0; + this.isStreaming = false; + this.ttftVal = 0; + this.durationVal = 0; + return this; + } + + @Override + public CallRecord provider(String provider) { + this.providerVal = provider; + return this; + } + + @Override + public CallRecord inputTokens(int tokens) { + this.inputTok = tokens; + return this; + } + + @Override + public CallRecord outputTokens(int tokens) { + this.outputTok = tokens; + return this; + } + + @Override + public CallRecord cacheReadTokens(int tokens) { + this.cacheReadTok = tokens; + return this; + } + + @Override + public CallRecord cacheCreationTokens(int tokens) { + this.cacheCreateTok = tokens; + return this; + } + + @Override + public CallRecord streaming() { + this.isStreaming = true; + return this; + } + + @Override + public CallRecord timeToFirstToken(long nanos) { + this.ttftVal = nanos; + return this; + } + + @Override + public CallRecord duration(long nanos) { + this.durationVal = nanos; + return this; + } + + @Override + public void record() { + service.commitCallRecord(this); + } +} diff --git a/btrace-extensions/btrace-llm-trace/src/main/java/io/btrace/llm/LlmTraceService.java b/btrace-extensions/btrace-llm-trace/src/main/java/io/btrace/llm/LlmTraceService.java new file mode 100644 index 000000000..3e8b20f3a --- /dev/null +++ b/btrace-extensions/btrace-llm-trace/src/main/java/io/btrace/llm/LlmTraceService.java @@ -0,0 +1,173 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.llm; + +import io.btrace.core.extensions.ServiceDescriptor; +import org.jetbrains.annotations.Nullable; + +/** + * LLM inference observability service for BTrace scripts. + * + *

Records LLM API call metrics — token counts, latencies, costs — and provides aggregated + * statistics. Scripts use {@code @OnMethod} to intercept LLM SDK calls and feed data into this + * service. + * + *

Thread-safe. All methods can be called concurrently from instrumented threads. + * + *

Quick start — latency only (works with any SDK)

+ * + * 
+ * @Injected private static LlmTraceService llm;
+ *
+ * @OnMethod(clazz = "+dev.langchain4j.model.chat.ChatLanguageModel",
+ *           method = "generate", location = @Location(Kind.RETURN))
+ * public static void onChat(@ProbeClassName String cls, @Duration long dur) {
+ *     llm.recordCall(cls, dur);
+ * }
+ *
+ * + *

Full metrics with fluent builder

+ * + * 
+ * llm.call("claude-sonnet-4-20250514")
+ *     .provider("anthropic")
+ *     .inputTokens(1500)
+ *     .outputTokens(300)
+ *     .cacheReadTokens(800)
+ *     .duration(durationNanos)
+ *     .record();
+ *
+ */ +@ServiceDescriptor +public interface LlmTraceService { + + // ==================== Simple recording ==================== + + /** + * Records an LLM call with only latency (no token info). Use this when token counts aren't easily + * extractable. + * + * @param model model identifier or class name + * @param durationNanos call duration in nanoseconds + */ + void recordCall(String model, long durationNanos); + + /** + * Records an LLM call with token counts. + * + * @param model model identifier (e.g. "gpt-4o", "claude-sonnet-4-20250514") + * @param inputTokens number of input/prompt tokens + * @param outputTokens number of output/completion tokens + * @param durationNanos call duration in nanoseconds + */ + void recordCall(String model, int inputTokens, int outputTokens, long durationNanos); + + // ==================== Fluent builder ==================== + + /** + * Starts a fluent call record for the given model. + * + *

The returned builder is allocation-free (ThreadLocal-pooled). It is safe to use on hot + * paths. The builder must be used inline on the calling thread and {@link CallRecord#record()} + * must be called before the next {@code call()}. + * + * 

+   * llm.call("claude-sonnet-4-20250514")
+   *     .provider("anthropic")
+   *     .inputTokens(1500)
+   *     .outputTokens(300)
+   *     .cacheReadTokens(800)
+   *     .streaming()
+   *     .timeToFirstToken(200_000_000L)
+   *     .duration(durationNanos)
+   *     .record();
+   *
+ * + * @param model model identifier + * @return a call record builder (thread-local, do not store) + */ + @Nullable CallRecord call(String model); + + // ==================== Specialized recording ==================== + + /** + * Records an embedding API call. + * + * @param model embedding model identifier (e.g. "text-embedding-3-small") + * @param tokenCount number of tokens embedded + * @param durationNanos call duration in nanoseconds + */ + void recordEmbedding(String model, int tokenCount, long durationNanos); + + /** + * Records a tool/function call invocation by the LLM. + * + * @param model model that made the tool call + * @param toolName name of the tool/function called + */ + void recordToolUse(String model, String toolName); + + /** + * Records a failed LLM API call. + * + * @param model model identifier + * @param errorType error class name or HTTP status code + * @param durationNanos call duration before failure + */ + void recordError(String model, String errorType, long durationNanos); + + // ==================== Reporting ==================== + + /** + * Returns a formatted summary of all recorded metrics. Includes per-model token counts, latency + * stats, cost estimates, cache hit rates, tool use, and error breakdown. + * + * @return multi-line summary string + */ + @Nullable String getSummary(); + + /** + * Returns a one-line summary for a specific model. + * + * @param model model identifier + * @return summary string, or "No data" if none recorded + */ + @Nullable String getModelSummary(String model); + + /** + * Returns total estimated cost across all models in USD. Uses built-in pricing table. Returns -1 + * if all models are unknown. + * + * @return estimated cost in USD, or -1 if pricing unknown + */ + double getEstimatedCostUsd(); + + /** Returns total number of chat completion calls recorded. */ + long getTotalCalls(); + + /** Returns total input tokens across all calls. */ + long getTotalInputTokens(); + + /** Returns total output tokens across all calls. */ + long getTotalOutputTokens(); + + /** Returns total number of embedding calls recorded. */ + long getTotalEmbeddingCalls(); + + /** Resets all collected metrics. */ + void reset(); +} diff --git a/btrace-extensions/btrace-llm-trace/src/main/java/io/btrace/llm/LlmTraceServiceImpl.java b/btrace-extensions/btrace-llm-trace/src/main/java/io/btrace/llm/LlmTraceServiceImpl.java new file mode 100644 index 000000000..3c9ed33ec --- /dev/null +++ b/btrace-extensions/btrace-llm-trace/src/main/java/io/btrace/llm/LlmTraceServiceImpl.java @@ -0,0 +1,487 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.llm; + +import io.btrace.core.extensions.Extension; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicLong; + +/** Thread-safe LLM call tracing with lock-free per-model statistics. Zero external dependencies. */ +public final class LlmTraceServiceImpl extends Extension implements LlmTraceService { + + private final Map modelStats = new ConcurrentHashMap<>(); + private final Map embeddingStats = new ConcurrentHashMap<>(); + private final Map toolUseCounts = new ConcurrentHashMap<>(); + private final Map errorCounts = new ConcurrentHashMap<>(); + + // ==================== Simple recording ==================== + + @Override + public void recordCall(String model, long durationNanos) { + ModelStats stats = getOrCreate(modelStats, model); + stats.calls.incrementAndGet(); + stats.totalDurationNanos.addAndGet(durationNanos); + updateMinMax(stats.minDurationNanos, stats.maxDurationNanos, durationNanos); + } + + @Override + public void recordCall(String model, int inputTokens, int outputTokens, long durationNanos) { + ModelStats stats = getOrCreate(modelStats, model); + stats.calls.incrementAndGet(); + stats.inputTokens.addAndGet(inputTokens); + stats.outputTokens.addAndGet(outputTokens); + stats.totalDurationNanos.addAndGet(durationNanos); + updateMinMax(stats.minDurationNanos, stats.maxDurationNanos, durationNanos); + } + + // ==================== Fluent builder ==================== + + /** + * ThreadLocal-pooled builder — one CallRecordImpl per thread, reused across calls. Eliminates + * per-call heap allocation, making the builder safe for hot paths. + */ + private final ThreadLocal callRecordPool = + ThreadLocal.withInitial(CallRecordImpl::new); + + @Override + public CallRecord call(String model) { + return callRecordPool.get().reset(this, model); + } + + void commitCallRecord(CallRecordImpl rec) { + ModelStats stats = getOrCreate(modelStats, rec.model); + if (rec.providerVal != null) { + stats.provider = rec.providerVal; + } + stats.calls.incrementAndGet(); + stats.inputTokens.addAndGet(rec.inputTok); + stats.outputTokens.addAndGet(rec.outputTok); + stats.cacheReadTokens.addAndGet(rec.cacheReadTok); + stats.cacheCreationTokens.addAndGet(rec.cacheCreateTok); + stats.totalDurationNanos.addAndGet(rec.durationVal); + updateMinMax(stats.minDurationNanos, stats.maxDurationNanos, rec.durationVal); + if (rec.isStreaming) { + stats.streamingCalls.incrementAndGet(); + stats.totalTimeToFirstToken.addAndGet(rec.ttftVal); + } + } + + // ==================== Specialized recording ==================== + + @Override + public void recordEmbedding(String model, int tokenCount, long durationNanos) { + ModelStats stats = getOrCreate(embeddingStats, model); + stats.calls.incrementAndGet(); + stats.inputTokens.addAndGet(tokenCount); + stats.totalDurationNanos.addAndGet(durationNanos); + updateMinMax(stats.minDurationNanos, stats.maxDurationNanos, durationNanos); + } + + @Override + public void recordToolUse(String model, String toolName) { + String key = model + "::" + toolName; + toolUseCounts.computeIfAbsent(key, k -> new AtomicLong()).incrementAndGet(); + getOrCreate(modelStats, model).toolCalls.incrementAndGet(); + } + + @Override + public void recordError(String model, String errorType, long durationNanos) { + String key = model + "::" + errorType; + errorCounts.computeIfAbsent(key, k -> new AtomicLong()).incrementAndGet(); + getOrCreate(modelStats, model).errors.incrementAndGet(); + } + + // ==================== Reporting ==================== + + @Override + public String getSummary() { + if (modelStats.isEmpty() && embeddingStats.isEmpty()) { + return "No LLM calls recorded."; + } + + StringBuilder sb = new StringBuilder(); + sb.append("=== LLM Trace Summary ===\n\n"); + + long totalCalls = 0; + long totalIn = 0; + long totalOut = 0; + double totalCost = 0; + + // Chat completions + for (Map.Entry entry : modelStats.entrySet()) { + String model = entry.getKey(); + ModelStats s = entry.getValue(); + long calls = s.calls.get(); + long inTok = s.inputTokens.get(); + long outTok = s.outputTokens.get(); + long cacheRead = s.cacheReadTokens.get(); + long cacheCreate = s.cacheCreationTokens.get(); + + totalCalls += calls; + totalIn += inTok; + totalOut += outTok; + + sb.append("Model: ").append(model); + if (!"unknown".equals(s.provider)) { + sb.append(" (").append(s.provider).append(")"); + } + sb.append("\n"); + + // Calls + sb.append(" Calls: ").append(calls); + long streaming = s.streamingCalls.get(); + if (streaming > 0) { + sb.append(" (").append(streaming).append(" streaming)"); + } + sb.append("\n"); + + // Tokens + if (inTok > 0 || outTok > 0) { + sb.append(" Tokens: ").append(inTok).append(" in / ").append(outTok).append(" out"); + if (calls > 0) { + sb.append(" (avg ").append(inTok / calls).append("/").append(outTok / calls).append(")"); + } + sb.append("\n"); + } + + // Cache + if (cacheRead > 0 || cacheCreate > 0) { + sb.append(" Cache: "); + if (cacheRead > 0) { + sb.append(cacheRead).append(" read"); + // Show cache hit rate relative to total input + if (inTok > 0) { + long hitPct = (cacheRead * 100) / inTok; + sb.append(" (").append(hitPct).append("% hit)"); + } + } + if (cacheCreate > 0) { + if (cacheRead > 0) sb.append(", "); + sb.append(cacheCreate).append(" created"); + } + sb.append("\n"); + } + + // Latency + if (calls > 0) { + long avgMs = (s.totalDurationNanos.get() / calls) / 1_000_000; + long minMs = s.minDurationNanos.get() / 1_000_000; + long maxMs = s.maxDurationNanos.get() / 1_000_000; + sb.append(" Latency: avg ").append(avgMs).append("ms"); + sb.append(", min ").append(minMs).append("ms"); + sb.append(", max ").append(maxMs).append("ms\n"); + } + + // TTFT + if (streaming > 0) { + long avgTtft = (s.totalTimeToFirstToken.get() / streaming) / 1_000_000; + sb.append(" TTFT (avg): ").append(avgTtft).append("ms\n"); + } + + // Tool calls + long tc = s.toolCalls.get(); + if (tc > 0) { + sb.append(" Tool calls: ").append(tc).append("\n"); + } + + // Errors + long errs = s.errors.get(); + if (errs > 0) { + sb.append(" Errors: ").append(errs).append("\n"); + } + + // Cost + double cost = estimateCost(model, inTok, outTok, cacheRead); + if (cost >= 0) { + totalCost += cost; + sb.append(" Est. cost: $").append(formatCost(cost)); + if (cacheRead > 0) { + double uncachedCost = estimateCost(model, inTok + cacheRead, outTok, 0); + if (uncachedCost > cost) { + sb.append(" (saved $").append(formatCost(uncachedCost - cost)).append(" via cache)"); + } + } + sb.append("\n"); + } + sb.append("\n"); + } + + // Embeddings + if (!embeddingStats.isEmpty()) { + sb.append("--- Embeddings ---\n"); + for (Map.Entry entry : embeddingStats.entrySet()) { + ModelStats s = entry.getValue(); + long calls = s.calls.get(); + long tokens = s.inputTokens.get(); + long avgMs = calls > 0 ? (s.totalDurationNanos.get() / calls) / 1_000_000 : 0; + sb.append(" ") + .append(entry.getKey()) + .append(": ") + .append(calls) + .append(" calls, ") + .append(tokens) + .append(" tokens, avg ") + .append(avgMs) + .append("ms\n"); + } + sb.append("\n"); + } + + // Tool use breakdown + if (!toolUseCounts.isEmpty()) { + sb.append("--- Tool Use ---\n"); + for (Map.Entry entry : toolUseCounts.entrySet()) { + sb.append(" ") + .append(entry.getKey()) + .append(": ") + .append(entry.getValue().get()) + .append("\n"); + } + sb.append("\n"); + } + + // Error breakdown + if (!errorCounts.isEmpty()) { + sb.append("--- Errors ---\n"); + for (Map.Entry entry : errorCounts.entrySet()) { + sb.append(" ") + .append(entry.getKey()) + .append(": ") + .append(entry.getValue().get()) + .append("\n"); + } + sb.append("\n"); + } + + sb.append("--- Totals ---\n"); + sb.append(" Calls: ").append(totalCalls).append("\n"); + if (totalIn > 0 || totalOut > 0) { + sb.append(" Tokens: ").append(totalIn).append(" in / ").append(totalOut).append(" out\n"); + } + if (totalCost > 0) { + sb.append(" Est. total cost: $").append(formatCost(totalCost)).append("\n"); + } + + return sb.toString(); + } + + @Override + public String getModelSummary(String model) { + ModelStats s = modelStats.get(model); + if (s == null) { + return "No data for model: " + model; + } + long calls = s.calls.get(); + long inTok = s.inputTokens.get(); + long outTok = s.outputTokens.get(); + long avgMs = calls > 0 ? (s.totalDurationNanos.get() / calls) / 1_000_000 : 0; + + StringBuilder sb = new StringBuilder(); + sb.append(model).append(": ").append(calls).append(" calls"); + if (inTok > 0 || outTok > 0) { + sb.append(", ").append(inTok).append("/").append(outTok).append(" tokens (in/out)"); + } + sb.append(", avg ").append(avgMs).append("ms"); + return sb.toString(); + } + + @Override + public double getEstimatedCostUsd() { + double total = 0; + boolean anyKnown = false; + for (Map.Entry entry : modelStats.entrySet()) { + ModelStats s = entry.getValue(); + double cost = + estimateCost( + entry.getKey(), s.inputTokens.get(), s.outputTokens.get(), s.cacheReadTokens.get()); + if (cost >= 0) { + total += cost; + anyKnown = true; + } + } + return anyKnown ? total : -1; + } + + @Override + public long getTotalCalls() { + long total = 0; + for (ModelStats s : modelStats.values()) { + total += s.calls.get(); + } + return total; + } + + @Override + public long getTotalInputTokens() { + long total = 0; + for (ModelStats s : modelStats.values()) { + total += s.inputTokens.get(); + } + return total; + } + + @Override + public long getTotalOutputTokens() { + long total = 0; + for (ModelStats s : modelStats.values()) { + total += s.outputTokens.get(); + } + return total; + } + + @Override + public long getTotalEmbeddingCalls() { + long total = 0; + for (ModelStats s : embeddingStats.values()) { + total += s.calls.get(); + } + return total; + } + + @Override + public void reset() { + modelStats.clear(); + embeddingStats.clear(); + toolUseCounts.clear(); + errorCounts.clear(); + } + + @Override + public void close() { + String summary = getSummary(); + if (!"No LLM calls recorded.".equals(summary)) { + getContext().send(summary); + } + } + + // ==================== Internals ==================== + + private static ModelStats getOrCreate(Map map, String key) { + return map.computeIfAbsent(key, k -> new ModelStats()); + } + + private static void updateMinMax(AtomicLong min, AtomicLong max, long value) { + long cur; + do { + cur = min.get(); + if (value >= cur) break; + } while (!min.compareAndSet(cur, value)); + do { + cur = max.get(); + if (value <= cur) break; + } while (!max.compareAndSet(cur, value)); + } + + /** + * Estimates cost in USD. Cache-read tokens are priced at ~10% of input rate for models that + * support caching. + */ + static double estimateCost( + String model, long inputTokens, long outputTokens, long cacheReadTokens) { + double inputPer1M = -1; + double outputPer1M = -1; + double cacheReadPer1M = -1; + + String m = model.toLowerCase(); + + // Anthropic Claude + if (m.contains("claude") && m.contains("opus")) { + inputPer1M = 15.0; + outputPer1M = 75.0; + cacheReadPer1M = 1.50; + } else if (m.contains("claude") && m.contains("sonnet")) { + inputPer1M = 3.0; + outputPer1M = 15.0; + cacheReadPer1M = 0.30; + } else if (m.contains("claude") && m.contains("haiku")) { + inputPer1M = 0.80; + outputPer1M = 4.0; + cacheReadPer1M = 0.08; + } + // OpenAI GPT + else if (m.contains("gpt-4o-mini")) { + inputPer1M = 0.15; + outputPer1M = 0.60; + cacheReadPer1M = 0.075; + } else if (m.contains("gpt-4o")) { + inputPer1M = 2.50; + outputPer1M = 10.0; + cacheReadPer1M = 1.25; + } else if (m.contains("gpt-4") && m.contains("turbo")) { + inputPer1M = 10.0; + outputPer1M = 30.0; + } else if (m.contains("gpt-4")) { + inputPer1M = 30.0; + outputPer1M = 60.0; + } else if (m.contains("gpt-3.5")) { + inputPer1M = 0.50; + outputPer1M = 1.50; + } else if (m.contains("o1-mini")) { + inputPer1M = 3.0; + outputPer1M = 12.0; + cacheReadPer1M = 1.50; + } else if (m.contains("o1")) { + inputPer1M = 15.0; + outputPer1M = 60.0; + cacheReadPer1M = 7.50; + } + // Google Gemini + else if (m.contains("gemini") && m.contains("pro")) { + inputPer1M = 1.25; + outputPer1M = 5.0; + } else if (m.contains("gemini") && m.contains("flash")) { + inputPer1M = 0.075; + outputPer1M = 0.30; + } + + if (inputPer1M < 0) { + return -1; + } + + double cost = + (inputTokens * inputPer1M / 1_000_000.0) + (outputTokens * outputPer1M / 1_000_000.0); + if (cacheReadTokens > 0 && cacheReadPer1M > 0) { + cost += (cacheReadTokens * cacheReadPer1M / 1_000_000.0); + } + return cost; + } + + static String formatCost(double cost) { + if (cost < 0.01) { + return String.format("%.4f", cost); + } + return String.format("%.2f", cost); + } + + /** Lock-free per-model statistics. */ + static final class ModelStats { + volatile String provider = "unknown"; + final AtomicLong calls = new AtomicLong(); + final AtomicLong streamingCalls = new AtomicLong(); + final AtomicLong inputTokens = new AtomicLong(); + final AtomicLong outputTokens = new AtomicLong(); + final AtomicLong cacheReadTokens = new AtomicLong(); + final AtomicLong cacheCreationTokens = new AtomicLong(); + final AtomicLong totalDurationNanos = new AtomicLong(); + final AtomicLong minDurationNanos = new AtomicLong(Long.MAX_VALUE); + final AtomicLong maxDurationNanos = new AtomicLong(0); + final AtomicLong totalTimeToFirstToken = new AtomicLong(); + final AtomicLong toolCalls = new AtomicLong(); + final AtomicLong errors = new AtomicLong(); + } +} diff --git a/btrace-extensions/btrace-llm-trace/src/test/java/io/btrace/llm/LlmTraceServiceTest.java b/btrace-extensions/btrace-llm-trace/src/test/java/io/btrace/llm/LlmTraceServiceTest.java new file mode 100644 index 000000000..81eff612b --- /dev/null +++ b/btrace-extensions/btrace-llm-trace/src/test/java/io/btrace/llm/LlmTraceServiceTest.java @@ -0,0 +1,356 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.llm; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class LlmTraceServiceTest { + + private LlmTraceServiceImpl service; + + @BeforeEach + void setUp() { + service = new LlmTraceServiceImpl(); + } + + // ==================== Simple recording ==================== + + @Test + void recordCallDurationOnly() { + service.recordCall("gpt-4o", 500_000_000L); + service.recordCall("gpt-4o", 300_000_000L); + + assertEquals(2, service.getTotalCalls()); + assertEquals(0, service.getTotalInputTokens(), "No tokens recorded"); + String summary = service.getModelSummary("gpt-4o"); + assertTrue(summary.contains("2 calls")); + assertTrue(summary.contains("400ms"), "avg of 500ms and 300ms"); + } + + @Test + void recordCallWithTokens() { + service.recordCall("claude-sonnet-4-20250514", 1000, 200, 500_000_000L); + service.recordCall("claude-sonnet-4-20250514", 800, 300, 400_000_000L); + + assertEquals(2, service.getTotalCalls()); + assertEquals(1800, service.getTotalInputTokens()); + assertEquals(500, service.getTotalOutputTokens()); + } + + // ==================== Fluent builder ==================== + + @Test + void fluentBuilderBasic() { + service.call("gpt-4o").inputTokens(1000).outputTokens(200).duration(500_000_000L).record(); + + assertEquals(1, service.getTotalCalls()); + assertEquals(1000, service.getTotalInputTokens()); + assertEquals(200, service.getTotalOutputTokens()); + } + + @Test + void fluentBuilderWithProvider() { + service + .call("claude-sonnet-4-20250514") + .provider("anthropic") + .inputTokens(1500) + .outputTokens(300) + .duration(800_000_000L) + .record(); + + String summary = service.getSummary(); + assertTrue(summary.contains("anthropic")); + assertTrue(summary.contains("claude-sonnet")); + } + + @Test + void fluentBuilderWithCache() { + service + .call("claude-sonnet-4-20250514") + .inputTokens(500) + .outputTokens(200) + .cacheReadTokens(1000) + .cacheCreationTokens(200) + .duration(300_000_000L) + .record(); + + String summary = service.getSummary(); + assertTrue(summary.contains("Cache:"), "Should show cache section"); + assertTrue(summary.contains("1000 read"), "Should show cache read tokens"); + assertTrue(summary.contains("200 created"), "Should show cache creation tokens"); + assertTrue(summary.contains("saved $"), "Should show cache savings"); + } + + @Test + void fluentBuilderStreaming() { + service + .call("gpt-4o") + .inputTokens(2000) + .outputTokens(500) + .streaming() + .timeToFirstToken(200_000_000L) + .duration(3_000_000_000L) + .record(); + + String summary = service.getSummary(); + assertTrue(summary.contains("streaming")); + assertTrue(summary.contains("TTFT")); + } + + @Test + void fluentBuilderDurationOnly() { + // Minimal usage — just model + duration + service.call("some-model").duration(100_000_000L).record(); + + assertEquals(1, service.getTotalCalls()); + } + + // ==================== Embeddings ==================== + + @Test + void recordEmbedding() { + service.recordEmbedding("text-embedding-3-small", 500, 50_000_000L); + service.recordEmbedding("text-embedding-3-small", 300, 30_000_000L); + + assertEquals(2, service.getTotalEmbeddingCalls()); + assertEquals(0, service.getTotalCalls(), "Embeddings don't count as chat calls"); + + String summary = service.getSummary(); + assertTrue(summary.contains("Embeddings")); + assertTrue(summary.contains("text-embedding-3-small")); + assertTrue(summary.contains("2 calls")); + assertTrue(summary.contains("800 tokens")); + } + + // ==================== Tool use ==================== + + @Test + void recordToolUse() { + service.recordCall("claude-sonnet-4-20250514", 1000, 200, 500_000_000L); + service.recordToolUse("claude-sonnet-4-20250514", "search_web"); + service.recordToolUse("claude-sonnet-4-20250514", "search_web"); + service.recordToolUse("claude-sonnet-4-20250514", "run_code"); + + String summary = service.getSummary(); + assertTrue(summary.contains("Tool calls: 3")); + assertTrue(summary.contains("Tool Use")); + assertTrue(summary.contains("search_web")); + } + + // ==================== Errors ==================== + + @Test + void recordError() { + service.recordCall("gpt-4o", 1000, 0, 100_000_000L); + service.recordError("gpt-4o", "RateLimitException", 50_000_000L); + + String summary = service.getSummary(); + assertTrue(summary.contains("Errors: 1")); + assertTrue(summary.contains("RateLimitException")); + } + + // ==================== Model summary ==================== + + @Test + void getModelSummaryUnknown() { + assertTrue(service.getModelSummary("nonexistent").contains("No data")); + } + + @Test + void getModelSummaryDurationOnly() { + service.recordCall("my-model", 500_000_000L); + String result = service.getModelSummary("my-model"); + assertTrue(result.contains("1 calls")); + assertTrue(result.contains("500ms")); + // Should NOT show "0/0 tokens" + assertFalse(result.contains("0/0")); + } + + @Test + void getModelSummaryWithTokens() { + service.recordCall("gpt-4o", 1000, 200, 500_000_000L); + String result = service.getModelSummary("gpt-4o"); + assertTrue(result.contains("1000/200 tokens")); + } + + // ==================== Cost estimation ==================== + + @Test + void estimateCostClaudeSonnet() { + // Sonnet: $3/1M input, $15/1M output + double cost = + LlmTraceServiceImpl.estimateCost("claude-sonnet-4-20250514", 1_000_000, 100_000, 0); + assertEquals(3.0 + 1.5, cost, 0.01); + } + + @Test + void estimateCostWithCacheRead() { + // Sonnet: $3/1M input, $15/1M output, $0.30/1M cache-read + double cost = + LlmTraceServiceImpl.estimateCost("claude-sonnet-4-20250514", 500_000, 100_000, 500_000); + double expected = (500_000 * 3.0 / 1e6) + (100_000 * 15.0 / 1e6) + (500_000 * 0.30 / 1e6); + assertEquals(expected, cost, 0.001); + } + + @Test + void estimateCostGpt4o() { + double cost = LlmTraceServiceImpl.estimateCost("gpt-4o", 1_000_000, 1_000_000, 0); + assertEquals(2.50 + 10.0, cost, 0.01); + } + + @Test + void estimateCostUnknownModel() { + assertEquals(-1, LlmTraceServiceImpl.estimateCost("my-custom-model", 1000, 1000, 0)); + } + + @Test + void estimatedCostUsdAcrossModels() { + service.recordCall("claude-sonnet-4-20250514", 1_000_000, 0, 1_000_000_000L); + service.recordCall("gpt-4o", 1_000_000, 0, 1_000_000_000L); + + double cost = service.getEstimatedCostUsd(); + assertEquals(5.50, cost, 0.01); + } + + @Test + void estimatedCostUsdUnknownModels() { + service.recordCall("my-custom-model", 1000, 500, 100_000_000L); + assertEquals(-1, service.getEstimatedCostUsd()); + } + + // ==================== Latency tracking ==================== + + @Test + void latencyMinMax() { + service.recordCall("gpt-4o", 100, 50, 100_000_000L); + service.recordCall("gpt-4o", 100, 50, 500_000_000L); + service.recordCall("gpt-4o", 100, 50, 200_000_000L); + + String summary = service.getSummary(); + assertTrue(summary.contains("min 100ms")); + assertTrue(summary.contains("max 500ms")); + } + + // ==================== Reset ==================== + + @Test + void resetClearsEverything() { + service.recordCall("gpt-4o", 1000, 200, 500_000_000L); + service.recordEmbedding("text-embedding-3-small", 100, 10_000_000L); + service.recordToolUse("gpt-4o", "search"); + service.recordError("gpt-4o", "Timeout", 100_000_000L); + + service.reset(); + + assertEquals(0, service.getTotalCalls()); + assertEquals(0, service.getTotalInputTokens()); + assertEquals(0, service.getTotalOutputTokens()); + assertEquals(0, service.getTotalEmbeddingCalls()); + assertEquals("No LLM calls recorded.", service.getSummary()); + } + + // ==================== Edge cases ==================== + + @Test + void summaryEmptyWhenNoCalls() { + assertEquals("No LLM calls recorded.", service.getSummary()); + } + + @Test + void summaryOmitsTokenLineWhenNoTokens() { + service.recordCall("my-model", 500_000_000L); + String summary = service.getSummary(); + // Should show the model and latency but not "0 in / 0 out" + assertTrue(summary.contains("my-model")); + assertFalse(summary.contains("0 in / 0 out")); + } + + @Test + void cacheHitRateCalculation() { + service + .call("claude-sonnet-4-20250514") + .inputTokens(200) + .outputTokens(100) + .cacheReadTokens(800) + .duration(100_000_000L) + .record(); + + String summary = service.getSummary(); + // 800 cache reads out of 200 input = 400% is wrong semantically, + // so check that the cache section exists + assertTrue(summary.contains("Cache:")); + assertTrue(summary.contains("800 read")); + } + + // ==================== Concurrency ==================== + + @Test + void concurrentRecording() throws InterruptedException { + int threads = 8; + int callsPerThread = 1000; + Thread[] workers = new Thread[threads]; + + for (int t = 0; t < threads; t++) { + workers[t] = + new Thread( + () -> { + for (int i = 0; i < callsPerThread; i++) { + service.recordCall("gpt-4o", 100, 50, 10_000_000L); + } + }); + workers[t].start(); + } + for (Thread w : workers) { + w.join(); + } + + assertEquals(threads * callsPerThread, service.getTotalCalls()); + assertEquals(threads * callsPerThread * 100L, service.getTotalInputTokens()); + } + + @Test + void concurrentBuilderRecording() throws InterruptedException { + int threads = 4; + int callsPerThread = 500; + Thread[] workers = new Thread[threads]; + + for (int t = 0; t < threads; t++) { + workers[t] = + new Thread( + () -> { + for (int i = 0; i < callsPerThread; i++) { + service + .call("claude-sonnet-4-20250514") + .provider("anthropic") + .inputTokens(100) + .outputTokens(50) + .duration(10_000_000L) + .record(); + } + }); + workers[t].start(); + } + for (Thread w : workers) { + w.join(); + } + + assertEquals(threads * callsPerThread, service.getTotalCalls()); + } +} diff --git a/btrace-extensions/btrace-rag-quality/build.gradle b/btrace-extensions/btrace-rag-quality/build.gradle new file mode 100644 index 000000000..5f632dad2 --- /dev/null +++ b/btrace-extensions/btrace-rag-quality/build.gradle @@ -0,0 +1,35 @@ +plugins { + id 'io.btrace.extension' + alias(libs.plugins.shadow) +} + +java { + sourceCompatibility = 8 + targetCompatibility = 8 +} + +compileJava { + javaCompiler = javaToolchains.compilerFor { + languageVersion.set(JavaLanguageVersion.of(11)) + } +} + +btraceExtension { + id = 'btrace-rag-quality' + name = 'BTrace RAG Quality' + description = 'RAG pipeline observability: retrieval latency, similarity scores, chunk quality for vector DB clients' + services = ['io.btrace.rag.RagQualityService'] +} + +dependencies { + apiCompileOnly 'org.jetbrains:annotations:26.1.0' + apiCompileOnly project(':btrace-core') + implCompileOnly 'org.jetbrains:annotations:26.1.0' + implImplementation project(':btrace-core') + + testImplementation libs.junit.jupiter +} + +test { + useJUnitPlatform() +} diff --git a/btrace-extensions/btrace-rag-quality/src/main/java/io/btrace/rag/QueryRecord.java b/btrace-extensions/btrace-rag-quality/src/main/java/io/btrace/rag/QueryRecord.java new file mode 100644 index 000000000..045eec86e --- /dev/null +++ b/btrace-extensions/btrace-rag-quality/src/main/java/io/btrace/rag/QueryRecord.java @@ -0,0 +1,60 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.rag; + +/** + * Fluent builder for recording a vector DB query with detailed metrics. + * + *

Obtain via {@link RagQualityService#query(String)}. All setters are optional. Call {@link + * #record()} to commit the metrics. + * + *

Allocation-free: Instances are pooled per-thread internally. The returned + * reference must not be stored or shared across threads. + * + * 

+ * rag.query("pinecone")
+ *     .resultCount(5)
+ *     .topScore(0.92f)
+ *     .lowScore(0.71f)
+ *     .embeddingDimension(1536)
+ *     .duration(durationNanos)
+ *     .record();
+ *
+ */ +public interface QueryRecord { + + /** Number of results/chunks returned by the query. */ + QueryRecord resultCount(int count); + + /** Highest similarity score in the result set (0.0 to 1.0). */ + QueryRecord topScore(float score); + + /** Lowest similarity score in the result set (0.0 to 1.0). */ + QueryRecord lowScore(float score); + + /** Embedding dimension used for the query vector. */ + QueryRecord embeddingDimension(int dimension); + + /** Total tokens across all returned chunks. */ + QueryRecord totalChunkTokens(int tokens); + + /** Query duration in nanoseconds. */ + QueryRecord duration(long nanos); + + /** Commits this query record to the service. */ + void record(); +} diff --git a/btrace-extensions/btrace-rag-quality/src/main/java/io/btrace/rag/QueryRecordImpl.java b/btrace-extensions/btrace-rag-quality/src/main/java/io/btrace/rag/QueryRecordImpl.java new file mode 100644 index 000000000..4f86752f9 --- /dev/null +++ b/btrace-extensions/btrace-rag-quality/src/main/java/io/btrace/rag/QueryRecordImpl.java @@ -0,0 +1,89 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.rag; + +/** + * ThreadLocal-pooled builder implementing the {@link QueryRecord} fluent API. One instance per + * thread, reused across calls — zero heap allocation. + */ +final class QueryRecordImpl implements QueryRecord { + + String source; + int resultCountVal; + float topScoreVal; + float lowScoreVal; + int embDimension; + int chunkTokensVal; + long durationVal; + + private RagQualityServiceImpl service; + + QueryRecordImpl() {} + + QueryRecordImpl reset(RagQualityServiceImpl service, String source) { + this.service = service; + this.source = source; + this.resultCountVal = 0; + this.topScoreVal = -1f; + this.lowScoreVal = -1f; + this.embDimension = 0; + this.chunkTokensVal = 0; + this.durationVal = 0; + return this; + } + + @Override + public QueryRecord resultCount(int count) { + this.resultCountVal = count; + return this; + } + + @Override + public QueryRecord topScore(float score) { + this.topScoreVal = score; + return this; + } + + @Override + public QueryRecord lowScore(float score) { + this.lowScoreVal = score; + return this; + } + + @Override + public QueryRecord embeddingDimension(int dimension) { + this.embDimension = dimension; + return this; + } + + @Override + public QueryRecord totalChunkTokens(int tokens) { + this.chunkTokensVal = tokens; + return this; + } + + @Override + public QueryRecord duration(long nanos) { + this.durationVal = nanos; + return this; + } + + @Override + public void record() { + service.commitQueryRecord(this); + } +} diff --git a/btrace-extensions/btrace-rag-quality/src/main/java/io/btrace/rag/RagQualityService.java b/btrace-extensions/btrace-rag-quality/src/main/java/io/btrace/rag/RagQualityService.java new file mode 100644 index 000000000..a056fb65a --- /dev/null +++ b/btrace-extensions/btrace-rag-quality/src/main/java/io/btrace/rag/RagQualityService.java @@ -0,0 +1,129 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.rag; + +import io.btrace.core.extensions.ServiceDescriptor; +import org.jetbrains.annotations.Nullable; + +/** + * BTrace extension service for RAG (Retrieval-Augmented Generation) pipeline observability. + * + *

Tracks vector database query performance, retrieval quality metrics, and end-to-end RAG + * pipeline latency. Works with any vector DB client (Pinecone, Milvus, Weaviate, Chroma, pgvector, + * Qdrant). + * + *

Usage in a BTrace script: + * + * 

+ * @Injected RagQualityService rag;
+ *
+ * @OnMethod(clazz = "io.pinecone.PineconeClient", method = "query")
+ * void onQuery(@Duration long dur) {
+ *     rag.recordQuery("pinecone", 10, dur);
+ * }
+ *
+ */ +@ServiceDescriptor +public interface RagQualityService { + + // ==================== Query recording ==================== + + /** + * Records a vector DB query with duration only. + * + * @param source vector DB or index name (e.g. "pinecone", "milvus-products") + * @param durationNanos query duration in nanoseconds + */ + void recordQuery(String source, long durationNanos); + + /** + * Records a vector DB query with result count and duration. + * + * @param source vector DB or index name + * @param resultCount number of results/chunks returned (top-K) + * @param durationNanos query duration in nanoseconds + */ + void recordQuery(String source, int resultCount, long durationNanos); + + /** + * Starts a detailed query record builder. Allocation-free (ThreadLocal-pooled). + * + *

Use inline on the calling thread and call {@link QueryRecord#record()} before the next + * {@code query()} call. Do not store the returned reference. + * + * 

+   * rag.query("pinecone")
+   *     .resultCount(5)
+   *     .topScore(0.92f)
+   *     .lowScore(0.71f)
+   *     .embeddingDimension(1536)
+   *     .duration(durationNanos)
+   *     .record();
+   *
+ * + * @param source vector DB or index name + * @return a query record builder (thread-local, do not store) + */ + @Nullable QueryRecord query(String source); + + // ==================== Pipeline recording ==================== + + /** + * Records an end-to-end RAG pipeline invocation (retrieve + generate). + * + * @param pipelineName pipeline identifier + * @param retrievalNanos time spent in retrieval phase + * @param generationNanos time spent in generation phase + */ + void recordPipeline(String pipelineName, long retrievalNanos, long generationNanos); + + /** + * Records a chunk that was retrieved and used in context. Useful for tracking context window + * utilization. + * + * @param source vector DB or index name + * @param chunkTokens approximate token count of the chunk + */ + void recordChunk(String source, int chunkTokens); + + /** + * Records a retrieval that returned no results (empty context). + * + * @param source vector DB or index name + */ + void recordEmptyRetrieval(String source); + + // ==================== Reporting ==================== + + /** Returns a formatted summary of all RAG quality metrics. */ + @Nullable String getSummary(); + + /** Returns summary for a specific source/index. */ + @Nullable String getSourceSummary(String source); + + /** Returns total number of queries recorded across all sources. */ + long getTotalQueries(); + + /** Returns total number of empty retrievals. */ + long getTotalEmptyRetrievals(); + + /** Returns average similarity score across all queries that reported scores. */ + float getAverageTopScore(); + + /** Resets all collected metrics. */ + void reset(); +} diff --git a/btrace-extensions/btrace-rag-quality/src/main/java/io/btrace/rag/RagQualityServiceImpl.java b/btrace-extensions/btrace-rag-quality/src/main/java/io/btrace/rag/RagQualityServiceImpl.java new file mode 100644 index 000000000..1f2970185 --- /dev/null +++ b/btrace-extensions/btrace-rag-quality/src/main/java/io/btrace/rag/RagQualityServiceImpl.java @@ -0,0 +1,349 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.rag; + +import io.btrace.core.extensions.Extension; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicLong; + +/** Thread-safe RAG quality tracking with lock-free per-source statistics. */ +public final class RagQualityServiceImpl extends Extension implements RagQualityService { + + private final Map sourceStats = new ConcurrentHashMap<>(); + private final Map pipelineStats = new ConcurrentHashMap<>(); + + private final ThreadLocal queryRecordPool = + ThreadLocal.withInitial(QueryRecordImpl::new); + + // ==================== Simple recording ==================== + + @Override + public void recordQuery(String source, long durationNanos) { + SourceStats stats = getOrCreate(source); + stats.queries.incrementAndGet(); + stats.totalDurationNanos.addAndGet(durationNanos); + updateMinMax(stats.minDurationNanos, stats.maxDurationNanos, durationNanos); + } + + @Override + public void recordQuery(String source, int resultCount, long durationNanos) { + SourceStats stats = getOrCreate(source); + stats.queries.incrementAndGet(); + stats.totalResults.addAndGet(resultCount); + stats.totalDurationNanos.addAndGet(durationNanos); + updateMinMax(stats.minDurationNanos, stats.maxDurationNanos, durationNanos); + if (resultCount == 0) { + stats.emptyRetrievals.incrementAndGet(); + } + } + + // ==================== Fluent builder ==================== + + @Override + public QueryRecord query(String source) { + return queryRecordPool.get().reset(this, source); + } + + void commitQueryRecord(QueryRecordImpl rec) { + SourceStats stats = getOrCreate(rec.source); + stats.queries.incrementAndGet(); + stats.totalResults.addAndGet(rec.resultCountVal); + stats.totalDurationNanos.addAndGet(rec.durationVal); + updateMinMax(stats.minDurationNanos, stats.maxDurationNanos, rec.durationVal); + + if (rec.resultCountVal == 0) { + stats.emptyRetrievals.incrementAndGet(); + } + + if (rec.topScoreVal >= 0) { + stats.scoredQueries.incrementAndGet(); + // Accumulate scores as fixed-point (multiply by 10000) to avoid floating point atomics + stats.totalTopScore.addAndGet((long) (rec.topScoreVal * 10000)); + updateMin(stats.minTopScore, (long) (rec.topScoreVal * 10000)); + updateMax(stats.maxTopScore, (long) (rec.topScoreVal * 10000)); + } + if (rec.lowScoreVal >= 0) { + stats.totalLowScore.addAndGet((long) (rec.lowScoreVal * 10000)); + } + if (rec.chunkTokensVal > 0) { + stats.totalChunkTokens.addAndGet(rec.chunkTokensVal); + } + if (rec.embDimension > 0) { + stats.lastEmbeddingDimension = rec.embDimension; + } + } + + // ==================== Pipeline recording ==================== + + @Override + public void recordPipeline(String pipelineName, long retrievalNanos, long generationNanos) { + PipelineStats ps = pipelineStats.computeIfAbsent(pipelineName, k -> new PipelineStats()); + ps.invocations.incrementAndGet(); + ps.totalRetrievalNanos.addAndGet(retrievalNanos); + ps.totalGenerationNanos.addAndGet(generationNanos); + } + + @Override + public void recordChunk(String source, int chunkTokens) { + SourceStats stats = getOrCreate(source); + stats.totalChunkTokens.addAndGet(chunkTokens); + stats.chunkCount.incrementAndGet(); + } + + @Override + public void recordEmptyRetrieval(String source) { + getOrCreate(source).emptyRetrievals.incrementAndGet(); + } + + // ==================== Reporting ==================== + + @Override + public String getSummary() { + if (sourceStats.isEmpty() && pipelineStats.isEmpty()) { + return "No RAG queries recorded."; + } + + StringBuilder sb = new StringBuilder(); + sb.append("=== RAG Quality Summary ===\n\n"); + + long totalQueries = 0; + long totalEmpty = 0; + + for (Map.Entry entry : sourceStats.entrySet()) { + String source = entry.getKey(); + SourceStats s = entry.getValue(); + long queries = s.queries.get(); + totalQueries += queries; + long empty = s.emptyRetrievals.get(); + totalEmpty += empty; + + sb.append("Source: ").append(source).append("\n"); + sb.append(" Queries: ").append(queries); + if (empty > 0) { + sb.append(" (").append(empty).append(" empty, ").append(empty * 100 / queries).append("%)"); + } + sb.append("\n"); + + // Results + long totalRes = s.totalResults.get(); + if (totalRes > 0 && queries > 0) { + sb.append(" Results: ") + .append(totalRes) + .append(" total (avg ") + .append(totalRes / queries) + .append("/query)\n"); + } + + // Similarity scores + long scored = s.scoredQueries.get(); + if (scored > 0) { + float avgTop = (s.totalTopScore.get() / (float) scored) / 10000f; + float minTop = s.minTopScore.get() / 10000f; + float maxTop = s.maxTopScore.get() / 10000f; + sb.append(" Similarity (top): avg ").append(String.format("%.3f", avgTop)); + sb.append(", min ").append(String.format("%.3f", minTop)); + sb.append(", max ").append(String.format("%.3f", maxTop)); + sb.append("\n"); + + long totalLow = s.totalLowScore.get(); + if (totalLow > 0) { + float avgLow = (totalLow / (float) scored) / 10000f; + float spread = avgTop - avgLow; + sb.append(" Similarity (low): avg ").append(String.format("%.3f", avgLow)); + sb.append(" (spread ").append(String.format("%.3f", spread)).append(")\n"); + } + } + + // Latency + if (queries > 0) { + long avgMs = (s.totalDurationNanos.get() / queries) / 1_000_000; + long minMs = s.minDurationNanos.get() / 1_000_000; + long maxMs = s.maxDurationNanos.get() / 1_000_000; + sb.append(" Latency: avg ").append(avgMs).append("ms"); + sb.append(", min ").append(minMs).append("ms"); + sb.append(", max ").append(maxMs).append("ms\n"); + } + + // Chunks + long chunks = s.chunkCount.get(); + long chunkTokens = s.totalChunkTokens.get(); + if (chunkTokens > 0) { + sb.append(" Context: ").append(chunkTokens).append(" tokens"); + if (chunks > 0) { + sb.append(" (") + .append(chunks) + .append(" chunks, avg ") + .append(chunkTokens / chunks) + .append(" tok/chunk)"); + } + sb.append("\n"); + } + + sb.append("\n"); + } + + // Pipelines + if (!pipelineStats.isEmpty()) { + sb.append("--- Pipelines ---\n"); + for (Map.Entry entry : pipelineStats.entrySet()) { + PipelineStats ps = entry.getValue(); + long inv = ps.invocations.get(); + long avgRetMs = inv > 0 ? (ps.totalRetrievalNanos.get() / inv) / 1_000_000 : 0; + long avgGenMs = inv > 0 ? (ps.totalGenerationNanos.get() / inv) / 1_000_000 : 0; + sb.append(" ") + .append(entry.getKey()) + .append(": ") + .append(inv) + .append(" invocations, avg retrieval ") + .append(avgRetMs) + .append("ms, avg generation ") + .append(avgGenMs) + .append("ms\n"); + } + sb.append("\n"); + } + + sb.append("--- Totals ---\n"); + sb.append(" Queries: ").append(totalQueries).append("\n"); + if (totalEmpty > 0) { + sb.append(" Empty retrievals: ").append(totalEmpty); + if (totalQueries > 0) { + sb.append(" (").append(totalEmpty * 100 / totalQueries).append("%)"); + } + sb.append("\n"); + } + + return sb.toString(); + } + + @Override + public String getSourceSummary(String source) { + SourceStats s = sourceStats.get(source); + if (s == null) { + return "No data for source: " + source; + } + long queries = s.queries.get(); + long avgMs = queries > 0 ? (s.totalDurationNanos.get() / queries) / 1_000_000 : 0; + long empty = s.emptyRetrievals.get(); + StringBuilder sb = new StringBuilder(); + sb.append(source) + .append(": ") + .append(queries) + .append(" queries, avg ") + .append(avgMs) + .append("ms"); + if (empty > 0) { + sb.append(", ").append(empty).append(" empty"); + } + return sb.toString(); + } + + @Override + public long getTotalQueries() { + long total = 0; + for (SourceStats s : sourceStats.values()) { + total += s.queries.get(); + } + return total; + } + + @Override + public long getTotalEmptyRetrievals() { + long total = 0; + for (SourceStats s : sourceStats.values()) { + total += s.emptyRetrievals.get(); + } + return total; + } + + @Override + public float getAverageTopScore() { + long totalScored = 0; + long totalScore = 0; + for (SourceStats s : sourceStats.values()) { + totalScored += s.scoredQueries.get(); + totalScore += s.totalTopScore.get(); + } + if (totalScored == 0) return -1f; + return (totalScore / (float) totalScored) / 10000f; + } + + @Override + public void reset() { + sourceStats.clear(); + pipelineStats.clear(); + } + + @Override + public void close() { + String summary = getSummary(); + if (!"No RAG queries recorded.".equals(summary)) { + getContext().send(summary); + } + } + + // ==================== Internals ==================== + + private SourceStats getOrCreate(String source) { + return sourceStats.computeIfAbsent(source, k -> new SourceStats()); + } + + private static void updateMinMax(AtomicLong min, AtomicLong max, long value) { + updateMin(min, value); + updateMax(max, value); + } + + private static void updateMin(AtomicLong min, long value) { + long cur; + do { + cur = min.get(); + if (value >= cur) break; + } while (!min.compareAndSet(cur, value)); + } + + private static void updateMax(AtomicLong max, long value) { + long cur; + do { + cur = max.get(); + if (value <= cur) break; + } while (!max.compareAndSet(cur, value)); + } + + static final class SourceStats { + final AtomicLong queries = new AtomicLong(); + final AtomicLong totalResults = new AtomicLong(); + final AtomicLong emptyRetrievals = new AtomicLong(); + final AtomicLong scoredQueries = new AtomicLong(); + final AtomicLong totalTopScore = new AtomicLong(); // fixed-point * 10000 + final AtomicLong minTopScore = new AtomicLong(Long.MAX_VALUE); + final AtomicLong maxTopScore = new AtomicLong(0); + final AtomicLong totalLowScore = new AtomicLong(); // fixed-point * 10000 + final AtomicLong totalDurationNanos = new AtomicLong(); + final AtomicLong minDurationNanos = new AtomicLong(Long.MAX_VALUE); + final AtomicLong maxDurationNanos = new AtomicLong(0); + final AtomicLong totalChunkTokens = new AtomicLong(); + final AtomicLong chunkCount = new AtomicLong(); + volatile int lastEmbeddingDimension; + } + + static final class PipelineStats { + final AtomicLong invocations = new AtomicLong(); + final AtomicLong totalRetrievalNanos = new AtomicLong(); + final AtomicLong totalGenerationNanos = new AtomicLong(); + } +} diff --git a/btrace-extensions/btrace-rag-quality/src/test/java/io/btrace/rag/RagQualityServiceTest.java b/btrace-extensions/btrace-rag-quality/src/test/java/io/btrace/rag/RagQualityServiceTest.java new file mode 100644 index 000000000..6a28bd60a --- /dev/null +++ b/btrace-extensions/btrace-rag-quality/src/test/java/io/btrace/rag/RagQualityServiceTest.java @@ -0,0 +1,226 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.rag; + +import static org.junit.jupiter.api.Assertions.*; + +import java.util.concurrent.CountDownLatch; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class RagQualityServiceTest { + + private RagQualityServiceImpl service; + + @BeforeEach + void setUp() { + service = new RagQualityServiceImpl(); + } + + @Test + void durationOnlyQuery() { + service.recordQuery("pinecone", 5_000_000L); + assertEquals(1, service.getTotalQueries()); + String summary = service.getSummary(); + assertTrue(summary.contains("pinecone")); + assertTrue(summary.contains("1")); + } + + @Test + void queryWithResultCount() { + service.recordQuery("milvus", 10, 8_000_000L); + assertEquals(1, service.getTotalQueries()); + assertEquals(0, service.getTotalEmptyRetrievals()); + assertTrue(service.getSummary().contains("10 total")); + } + + @Test + void emptyRetrieval() { + service.recordQuery("weaviate", 0, 2_000_000L); + assertEquals(1, service.getTotalEmptyRetrievals()); + assertTrue(service.getSummary().contains("1 empty")); + } + + @Test + void explicitEmptyRetrieval() { + service.recordEmptyRetrieval("chroma"); + assertEquals(1, service.getTotalEmptyRetrievals()); + } + + @Test + void fluentBuilder() { + service + .query("pinecone") + .resultCount(5) + .topScore(0.92f) + .lowScore(0.71f) + .embeddingDimension(1536) + .totalChunkTokens(3000) + .duration(10_000_000L) + .record(); + + assertEquals(1, service.getTotalQueries()); + float avgScore = service.getAverageTopScore(); + assertTrue(avgScore > 0.91f && avgScore < 0.93f, "avgScore=" + avgScore); + String summary = service.getSummary(); + assertTrue(summary.contains("0.920")); + assertTrue(summary.contains("3000 tokens")); + } + + @Test + void fluentBuilderMinimal() { + service.query("qdrant").duration(1_000_000L).record(); + + assertEquals(1, service.getTotalQueries()); + assertEquals(-1f, service.getAverageTopScore()); + } + + @Test + void fluentBuilderEmptyResult() { + service.query("pgvector").resultCount(0).duration(500_000L).record(); + + assertEquals(1, service.getTotalEmptyRetrievals()); + } + + @Test + void multipleSourcesTrackedSeparately() { + service.recordQuery("pinecone", 5, 10_000_000L); + service.recordQuery("milvus", 3, 8_000_000L); + service.recordQuery("pinecone", 7, 12_000_000L); + + assertEquals(3, service.getTotalQueries()); + assertTrue(service.getSourceSummary("pinecone").contains("2 queries")); + assertTrue(service.getSourceSummary("milvus").contains("1 queries")); + } + + @Test + void pipelineRecording() { + service.recordPipeline("qa-bot", 5_000_000L, 50_000_000L); + service.recordPipeline("qa-bot", 3_000_000L, 45_000_000L); + + String summary = service.getSummary(); + assertTrue(summary.contains("qa-bot")); + assertTrue(summary.contains("2 invocations")); + } + + @Test + void chunkRecording() { + service.recordChunk("pinecone", 500); + service.recordChunk("pinecone", 750); + + String summary = service.getSourceSummary("pinecone"); + // Chunks don't count as queries + assertTrue(summary.contains("0 queries")); + } + + @Test + void similarityScoreAggregation() { + service.query("pinecone").topScore(0.90f).duration(1_000_000L).record(); + service.query("pinecone").topScore(0.80f).duration(1_000_000L).record(); + service.query("milvus").topScore(0.70f).duration(1_000_000L).record(); + + float avg = service.getAverageTopScore(); + // (0.90 + 0.80 + 0.70) / 3 = 0.8 + assertTrue(avg > 0.79f && avg < 0.81f, "avg=" + avg); + } + + @Test + void latencyMinMax() { + service.recordQuery("pinecone", 10_000_000L); + service.recordQuery("pinecone", 2_000_000L); + service.recordQuery("pinecone", 50_000_000L); + + String summary = service.getSummary(); + assertTrue(summary.contains("min 2ms")); + assertTrue(summary.contains("max 50ms")); + } + + @Test + void noDataSummary() { + assertEquals("No RAG queries recorded.", service.getSummary()); + } + + @Test + void unknownSourceSummary() { + assertEquals("No data for source: unknown", service.getSourceSummary("unknown")); + } + + @Test + void reset() { + service.recordQuery("pinecone", 5_000_000L); + service.recordEmptyRetrieval("milvus"); + service.recordPipeline("qa", 1L, 1L); + service.reset(); + + assertEquals(0, service.getTotalQueries()); + assertEquals(0, service.getTotalEmptyRetrievals()); + assertEquals("No RAG queries recorded.", service.getSummary()); + } + + @Test + void concurrentRecording() throws Exception { + int threads = 8; + int queriesPerThread = 1000; + CountDownLatch latch = new CountDownLatch(threads); + + for (int t = 0; t < threads; t++) { + new Thread( + () -> { + try { + for (int i = 0; i < queriesPerThread; i++) { + service.recordQuery("pinecone", 5, 1_000_000L); + } + } finally { + latch.countDown(); + } + }) + .start(); + } + latch.await(); + + assertEquals(threads * queriesPerThread, service.getTotalQueries()); + } + + @Test + void concurrentBuilderRecording() throws Exception { + int threads = 8; + int queriesPerThread = 500; + CountDownLatch latch = new CountDownLatch(threads); + + for (int t = 0; t < threads; t++) { + new Thread( + () -> { + try { + for (int i = 0; i < queriesPerThread; i++) { + service + .query("milvus") + .resultCount(3) + .topScore(0.85f) + .duration(2_000_000L) + .record(); + } + } finally { + latch.countDown(); + } + }) + .start(); + } + latch.await(); + + assertEquals(threads * queriesPerThread, service.getTotalQueries()); + } +} diff --git a/btrace-gradle-plugin/src/main/groovy/io/btrace/gradle/BTraceExtensionPlugin.groovy b/btrace-gradle-plugin/src/main/groovy/io/btrace/gradle/BTraceExtensionPlugin.groovy index 591a63686..f36619825 100644 --- a/btrace-gradle-plugin/src/main/groovy/io/btrace/gradle/BTraceExtensionPlugin.groovy +++ b/btrace-gradle-plugin/src/main/groovy/io/btrace/gradle/BTraceExtensionPlugin.groovy @@ -34,19 +34,19 @@ class BTraceExtensionPlugin implements Plugin { // Try to ensure Shadow is available; if resolution is blocked, we will emit a clear // error later with guidance. This is best-effort and safe when already applied. try { - if (!project.pluginManager.hasPlugin('com.github.johnrengelman.shadow')) { + if (!project.pluginManager.hasPlugin('com.gradleup.shadow')) { // Respect opt-out def ext = project.extensions.findByType(BTraceExtensionMetadata) boolean shouldAutoApply = (ext == null) ? true : (ext.autoApplyShadow != false) if (shouldAutoApply) { - project.logger.lifecycle("[BTRACE-EXT] Applying Shadow plugin automatically (com.github.johnrengelman.shadow) for ${project.path}") - project.pluginManager.apply('com.github.johnrengelman.shadow') + project.logger.lifecycle("[BTRACE-EXT] Applying Shadow plugin automatically (com.gradleup.shadow) for ${project.path}") + project.pluginManager.apply('com.gradleup.shadow') } else { project.logger.lifecycle("[BTRACE-EXT] Shadow auto-apply disabled (btraceExtension.autoApplyShadow=false) for ${project.path}") } } } catch (Throwable t) { - project.logger.warn("[BTRACE-EXT] Unable to auto-apply Shadow plugin: ${t.message}. Apply it explicitly via plugins { id 'com.github.johnrengelman.shadow' } or alias(libs.plugins.shadow), or set btraceExtension.autoApplyShadow=true.") + project.logger.warn("[BTRACE-EXT] Unable to auto-apply Shadow plugin: ${t.message}. Apply it explicitly via plugins { id 'com.gradleup.shadow' } or alias(libs.plugins.shadow), or set btraceExtension.autoApplyShadow=true.") } // Create extension for metadata @@ -604,7 +604,7 @@ class BTraceExtensionPlugin implements Plugin { // Task: Build Implementation JAR (requires Shadow plugin applied by the consumer project) def implJarProviderRef = new Object[1] - project.pluginManager.withPlugin('com.github.johnrengelman.shadow') { + project.pluginManager.withPlugin('com.gradleup.shadow') { def shadowJarProvider = project.tasks.named('shadowJar', Jar) project.afterEvaluate { shadowJarProvider.configure { @@ -634,7 +634,7 @@ class BTraceExtensionPlugin implements Plugin { def buildImplJar = project.tasks.register('buildImplJar') { doFirst { if (implJarProviderRef[0] == null) { - throw new IllegalStateException("Shadow plugin ('com.github.johnrengelman.shadow') must be applied in the project using the BTrace extension plugin.") + throw new IllegalStateException("Shadow plugin ('com.gradleup.shadow') must be applied in the project using the BTrace extension plugin.") } } dependsOn { implJarProviderRef[0] } @@ -646,7 +646,7 @@ class BTraceExtensionPlugin implements Plugin { def ext = project.extensions.findByType(BTraceExtensionMetadata) boolean autoApplied = (ext == null) ? true : (ext.autoApplyShadow != false) String hint = autoApplied ? "Ensure the Shadow plugin is resolvable/available." : "Enable auto-apply (btraceExtension.autoApplyShadow=true) or apply Shadow explicitly." - throw new GradleException("[BTRACE-EXT] Shadow plugin is required for ${project.path}. Apply id 'com.github.johnrengelman.shadow' (or alias(libs.plugins.shadow)). ${hint}") + throw new GradleException("[BTRACE-EXT] Shadow plugin is required for ${project.path}. Apply id 'com.gradleup.shadow' (or alias(libs.plugins.shadow)). ${hint}") } def implArchiveProvider = ((org.gradle.api.tasks.TaskProvider) implJarProviderRef[0]).flatMap { it.archiveFile } buildApiJar.configure { @@ -1371,6 +1371,7 @@ class BTraceExtensionPlugin implements Plugin { project.tasks.matching { it.name == 'check' }.configureEach { it.dependsOn validateServiceApis } } } + } class BTraceExtensionMetadata { diff --git a/btrace-gradle-plugin/src/main/groovy/io/btrace/gradle/BTraceFatAgentExtension.groovy b/btrace-gradle-plugin/src/main/groovy/io/btrace/gradle/BTraceFatAgentExtension.groovy index 31a93e430..1a06306da 100644 --- a/btrace-gradle-plugin/src/main/groovy/io/btrace/gradle/BTraceFatAgentExtension.groovy +++ b/btrace-gradle-plugin/src/main/groovy/io/btrace/gradle/BTraceFatAgentExtension.groovy @@ -25,6 +25,8 @@ import org.gradle.api.file.FileCollection * */ class BTraceFatAgentExtension { + static final String DEFAULT_REGISTRY_URL = 'https://btraceio.github.io/btrace-extensions/registry/extensions.json' + private final Project project private final List extensionSources = [] private final ProbeBundleSpec probeBundle @@ -53,9 +55,16 @@ class BTraceFatAgentExtension { /** Property name for filtering extensions when autoDiscover is true */ String filterProperty = 'embedExtensions' + /** Extension registry URL for resolving registry("id") sources. */ + String registryUrl = System.getProperty('btrace.extensions.registry', DEFAULT_REGISTRY_URL) + + /** Local cache file for the extension registry document. */ + File registryCacheFile + BTraceFatAgentExtension(Project project) { this.project = project this.outputDir = project.layout.buildDirectory.dir('libs').get().asFile + this.registryCacheFile = project.layout.buildDirectory.file('registry/extensions.json').get().asFile this.probeBundle = new ProbeBundleSpec(project) } @@ -63,7 +72,7 @@ class BTraceFatAgentExtension { * Configure embedded extensions. */ void embedExtensions(Action action) { - def spec = new ExtensionSourceSpec(project) + def spec = new ExtensionSourceSpec(project, this) action.execute(spec) extensionSources.addAll(spec.sources) } @@ -119,10 +128,12 @@ class BTraceFatAgentExtension { */ class ExtensionSourceSpec { private final Project project + private final BTraceFatAgentExtension owner final List sources = [] - ExtensionSourceSpec(Project project) { + ExtensionSourceSpec(Project project, BTraceFatAgentExtension owner) { this.project = project + this.owner = owner } /** diff --git a/btrace-gradle-plugin/src/test/java/io/btrace/gradle/BTraceExtensionPluginTest.java b/btrace-gradle-plugin/src/test/java/io/btrace/gradle/BTraceExtensionPluginTest.java index 36b4b1131..c59f09760 100644 --- a/btrace-gradle-plugin/src/test/java/io/btrace/gradle/BTraceExtensionPluginTest.java +++ b/btrace-gradle-plugin/src/test/java/io/btrace/gradle/BTraceExtensionPluginTest.java @@ -7,6 +7,7 @@ import java.io.File; import java.io.IOException; +import java.io.InputStream; import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.Path; @@ -154,6 +155,111 @@ void apiPublicationsStayApiOnly() throws IOException { } } + @Test + @DisplayName("updateRegistryCatalog writes entry into local registry checkout") + void updateRegistryCatalogWritesLocalRegistry() throws IOException { + Path registryDir = projectDir.resolve("registry-repo"); + Files.createDirectories(registryDir.resolve("registry")); + writeFile( + registryDir.resolve("registry/extensions.json"), + "{\n" + + " \"schema_version\": 1,\n" + + " \"extensions\": []\n" + + "}\n"); + + writeExtensionProject( + "btraceRegistry {\n" + + " prMode = 'off'\n" + + " verifyPublishedCoordinates = false\n" + + " registryWorktreeDir = file('" + + registryDir.toString().replace("\\", "/") + + "')\n" + + " tags = ['metrics']\n" + + "}\n"); + + BuildResult result = createRunner().withArguments(":ext:updateRegistryCatalog").build(); + + assertEquals(TaskOutcome.SUCCESS, result.task(":ext:updateRegistryCatalog").getOutcome()); + String json = + Files.readString( + registryDir.resolve("registry/extensions.json"), StandardCharsets.UTF_8); + assertTrue(json.contains("\"id\" : \"test.ext\"")); + assertTrue(json.contains("\"artifactId\" : \"ext\"")); + assertTrue(json.contains("\"version\" : \"1.0\"")); + } + + @Test + @DisplayName("updateRegistryCatalog uses a fork-based PR when no direct push repo is configured") + void updateRegistryCatalogUsesForkBasedPr() throws IOException { + Path registryDir = projectDir.resolve("registry-repo"); + Files.createDirectories(registryDir.resolve("registry")); + writeFile( + registryDir.resolve("registry/extensions.json"), + "{\n" + + " \"schema_version\": 1,\n" + + " \"extensions\": []\n" + + "}\n"); + initGitRepo(registryDir); + + Path forkRemote = projectDir.resolve("registry-fork.git"); + runCommand(projectDir, "git", "init", "--bare", forkRemote.toString()); + + Path ghScript = projectDir.resolve("fake-gh.sh"); + Path prLog = projectDir.resolve("gh-pr.log"); + writeFile( + ghScript, + "#!/bin/sh\n" + + "set -eu\n" + + "if [ \"$1\" = \"repo\" ] && [ \"$2\" = \"fork\" ]; then\n" + + " if git remote get-url registry-fork >/dev/null 2>&1; then\n" + + " git remote set-url registry-fork '" + forkRemote.toString().replace("\\", "/") + "'\n" + + " else\n" + + " git remote add registry-fork '" + forkRemote.toString().replace("\\", "/") + "'\n" + + " fi\n" + + " exit 0\n" + + "fi\n" + + "if [ \"$1\" = \"api\" ] && [ \"$2\" = \"user\" ]; then\n" + + " echo 'fork-user'\n" + + " exit 0\n" + + "fi\n" + + "if [ \"$1\" = \"pr\" ] && [ \"$2\" = \"create\" ]; then\n" + + " printf '%s\\n' \"$@\" > '" + prLog.toString().replace("\\", "/") + "'\n" + + " exit 0\n" + + "fi\n" + + "echo 'unexpected gh args: '$* >&2\n" + + "exit 1\n"); + ghScript.toFile().setExecutable(true); + + writeExtensionProject( + "btraceRegistry {\n" + + " prMode = 'auto'\n" + + " verifyPublishedCoordinates = false\n" + + " registryWorktreeDir = file('" + + registryDir.toString().replace("\\", "/") + + "')\n" + + " githubCli = '" + + ghScript.toString().replace("\\", "/") + + "'\n" + + " tags = ['metrics']\n" + + "}\n"); + + BuildResult result = createRunner().withArguments(":ext:updateRegistryCatalog").build(); + + assertEquals(TaskOutcome.SUCCESS, result.task(":ext:updateRegistryCatalog").getOutcome()); + String json = + Files.readString( + registryDir.resolve("registry/extensions.json"), StandardCharsets.UTF_8); + assertTrue(json.contains("\"id\" : \"test.ext\"")); + assertTrue(json.contains("\"artifactId\" : \"ext\"")); + assertTrue(json.contains("\"version\" : \"1.0\"")); + + String prArgs = Files.readString(prLog, StandardCharsets.UTF_8); + assertTrue(prArgs.contains("--repo")); + assertTrue(prArgs.contains("btraceio/btrace-extensions")); + assertTrue(prArgs.contains("--head")); + assertTrue(prArgs.contains("fork-user:btrace-registry/ext-1.0")); + } + private void writeStubCoreProject() throws IOException { Path dir = projectDir.resolve("btrace-core"); Files.createDirectories(dir.resolve("src/main/java/org/openjdk/btrace/core/extensions")); @@ -219,6 +325,10 @@ private void writeStubProcessorProject() throws IOException { } private void writeExtensionProject() throws IOException { + writeExtensionProject(""); + } + + private void writeExtensionProject(String extraBuildLogic) throws IOException { Path dir = projectDir.resolve("ext"); Files.createDirectories(dir.resolve("src/main/java/com/example/api")); Files.createDirectories(dir.resolve("src/main/java/com/example/impl")); @@ -228,7 +338,7 @@ private void writeExtensionProject() throws IOException { "plugins {\n" + " id 'java-library'\n" + " id 'io.btrace.extension'\n" - + " id 'com.github.johnrengelman.shadow'\n" + + " id 'com.gradleup.shadow'\n" + "}\n" + "group = 'com.example'\n" + "version = '1.0'\n" @@ -242,7 +352,8 @@ private void writeExtensionProject() throws IOException { + " services = ['com.example.api.PublicService']\n" + " requiredPermissions = ['NONE']\n" + " scanPermissions = false\n" - + "}\n"); + + "}\n" + + extraBuildLogic); writeFile( dir.resolve("src/main/java/com/example/api/PublicValue.java"), "package com.example.api;\n" @@ -322,6 +433,10 @@ private void writeBuildSrcShadowPlugin() throws IOException { projectDir.resolve( "buildSrc/src/main/resources/META-INF/gradle-plugins/com.github.johnrengelman.shadow.properties"), "implementation-class=com.github.jengelman.gradle.plugins.shadow.ShadowPlugin\n"); + writeFile( + projectDir.resolve( + "buildSrc/src/main/resources/META-INF/gradle-plugins/com.gradleup.shadow.properties"), + "implementation-class=com.github.jengelman.gradle.plugins.shadow.ShadowPlugin\n"); } private GradleRunner createRunner() { @@ -332,6 +447,33 @@ private GradleRunner createRunner() { .forwardOutput(); } + private void initGitRepo(Path dir) throws IOException { + runCommand(dir, "git", "init"); + runCommand(dir, "git", "-C", dir.toString(), "config", "user.name", "Test User"); + runCommand(dir, "git", "-C", dir.toString(), "config", "user.email", "test@example.com"); + runCommand(dir, "git", "-C", dir.toString(), "add", "."); + runCommand(dir, "git", "-C", dir.toString(), "commit", "-m", "init"); + } + + private void runCommand(Path dir, String... command) throws IOException { + ProcessBuilder builder = new ProcessBuilder(command); + builder.directory(dir.toFile()); + builder.redirectErrorStream(true); + Process process = builder.start(); + String output; + try (InputStream stream = process.getInputStream()) { + output = new String(stream.readAllBytes(), StandardCharsets.UTF_8); + } + try { + if (process.waitFor() != 0) { + throw new IOException("Command failed: " + String.join(" ", command) + "\n" + output); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted while running: " + String.join(" ", command), e); + } + } + private void writeFile(Path path, String content) throws IOException { Files.createDirectories(path.getParent()); Files.writeString(path, content, StandardCharsets.UTF_8); diff --git a/btrace-gradle-plugin/src/test/java/io/btrace/gradle/BTraceFatAgentPluginTest.java b/btrace-gradle-plugin/src/test/java/io/btrace/gradle/BTraceFatAgentPluginTest.java index b886e5e15..8723633eb 100644 --- a/btrace-gradle-plugin/src/test/java/io/btrace/gradle/BTraceFatAgentPluginTest.java +++ b/btrace-gradle-plugin/src/test/java/io/btrace/gradle/BTraceFatAgentPluginTest.java @@ -26,6 +26,7 @@ import java.io.File; import java.io.FileWriter; import java.io.IOException; +import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.Path; import org.gradle.testkit.runner.BuildResult; @@ -293,6 +294,38 @@ void stagingDirectoryCreated() throws IOException { assertTrue(Files.exists(stagingDir), "Staging directory should be created"); } + @Test + @DisplayName("Registry source fails clearly for unknown extension id") + void registrySourceFailsForUnknownId() throws IOException { + Path registry = projectDir.resolve("extensions.json"); + Files.writeString( + registry, + "{\n" + + " \"schema_version\": 1,\n" + + " \"extensions\": []\n" + + "}\n", + StandardCharsets.UTF_8); + + writeFile(buildFile, + "plugins {\n" + + " id 'io.btrace.fat-agent'\n" + + "}\n" + + "\n" + + "btraceFatAgent {\n" + + " registryUrl = '" + registry.toUri().toString() + "'\n" + + " embedExtensions {\n" + + " registry('missing-ext')\n" + + " }\n" + + "}\n" + ); + + BuildResult result = createRunner() + .withArguments("stageExtensions") + .buildAndFail(); + + assertTrue(result.getOutput().contains("Unknown extension id: missing-ext")); + } + private GradleRunner createRunner() { return GradleRunner.create() .withProjectDir(projectDir.toFile()) diff --git a/btrace-mcp-server/README.md b/btrace-mcp-server/README.md new file mode 100644 index 000000000..da2f5ca16 --- /dev/null +++ b/btrace-mcp-server/README.md @@ -0,0 +1,114 @@ +# BTrace MCP Server + +An MCP (Model Context Protocol) server that exposes BTrace operations as tools, allowing LLM clients (Claude Desktop, Claude Code, Cursor, etc.) to diagnose running JVMs. + +## Tools + +| Tool | Description | +|------|-------------| +| `list_jvms` | List all attachable Java VMs on this host | +| `deploy_oneliner` | Deploy a BTrace oneliner probe (e.g. `com.example.Service::method @return { print method, duration }`) | +| `deploy_script` | Deploy a full BTrace Java script | +| `list_probes` | List active probes on a JVM | +| `send_event` | Send an event to a running probe | +| `detach_probe` | Detach from a running probe (probe continues) | +| `exit_probe` | Stop and remove a probe | + +## Prompts + +| Prompt | Description | +|--------|-------------| +| `diagnose_slow_endpoint` | Step-by-step guide to diagnose a slow HTTP endpoint | +| `find_exception_source` | Guide to find where exceptions originate | +| `profile_method` | Guide to profile a method's latency | + +## Building + +```bash +./gradlew :btrace-mcp-server:build +``` + +## Configuration + +### Claude Desktop + +Add to `~/Library/Application Support/Claude/claude_desktop_config.json` (macOS) or `%APPDATA%\Claude\claude_desktop_config.json` (Windows): + +```json +{ + "mcpServers": { + "btrace": { + "command": "java", + "args": [ + "-cp", "/path/to/btrace/lib/*", + "io.btrace.mcp.BTraceMcpServer" + ] + } + } +} +``` + +### Claude Code + +Add to your project's `.mcp.json`: + +```json +{ + "mcpServers": { + "btrace": { + "command": "java", + "args": [ + "-cp", "/path/to/btrace/lib/*", + "io.btrace.mcp.BTraceMcpServer" + ] + } + } +} +``` + +Or configure globally in `~/.claude/settings.json`: + +```json +{ + "mcpServers": { + "btrace": { + "command": "java", + "args": [ + "-cp", "/path/to/btrace/lib/*", + "io.btrace.mcp.BTraceMcpServer" + ] + } + } +} +``` + +### Using the BTrace distribution + +If you have BTrace installed (e.g. via SDKMAN), you can reference the distribution directly: + +```json +{ + "mcpServers": { + "btrace": { + "command": "java", + "args": [ + "-cp", "$BTRACE_HOME/lib/*:$BTRACE_HOME/build/btrace-mcp-server.jar", + "io.btrace.mcp.BTraceMcpServer" + ] + } + } +} +``` + +## Usage Examples + +Once configured, you can ask your LLM client: + +- "List all running Java processes" +- "Attach to PID 12345 and trace all methods in com.example.UserService that take longer than 100ms" +- "Find where NullPointerExceptions are being thrown in PID 12345" +- "Profile the latency of com.example.OrderService::processOrder" + +## Protocol + +The server uses MCP over stdio (stdin/stdout JSON-RPC). All logging goes to stderr. diff --git a/btrace-mcp-server/build.gradle b/btrace-mcp-server/build.gradle new file mode 100644 index 000000000..3e13bdd6e --- /dev/null +++ b/btrace-mcp-server/build.gradle @@ -0,0 +1,28 @@ +// MCP server requires Java 11+ (uses ConcurrentHashMap.forEach, etc.) +compileJava { + sourceCompatibility = 11 + targetCompatibility = 11 +} + +dependencies { + implementation libs.slf4j + implementation libs.slf4j.simple + + def toolsJar = getToolsJar() + if (toolsJar.getAsFile().exists()) { + compileOnly files("${toolsJar}") + } + + implementation project(':btrace-core') + implementation project(':btrace-client') + implementation project(':btrace-compiler') + implementation project(':btrace-boot') +} + +jar { + manifest { + attributes( + 'Main-Class': 'io.btrace.mcp.BTraceMcpServer' + ) + } +} diff --git a/btrace-mcp-server/src/main/java/io/btrace/mcp/BTraceMcpServer.java b/btrace-mcp-server/src/main/java/io/btrace/mcp/BTraceMcpServer.java new file mode 100644 index 000000000..a7600cf57 --- /dev/null +++ b/btrace-mcp-server/src/main/java/io/btrace/mcp/BTraceMcpServer.java @@ -0,0 +1,217 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.mcp; + +import io.btrace.mcp.prompts.DiagnosticPrompts; +import io.btrace.mcp.tools.DeployOnelinerHandler; +import io.btrace.mcp.tools.DeployScriptHandler; +import io.btrace.mcp.tools.DetachProbeHandler; +import io.btrace.mcp.tools.ExitProbeHandler; +import io.btrace.mcp.tools.ListJvmsHandler; +import io.btrace.mcp.tools.ListProbesHandler; +import io.btrace.mcp.tools.SendEventHandler; +import java.io.IOException; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * BTrace MCP (Model Context Protocol) server. Exposes BTrace operations as MCP tools over stdio + * JSON-RPC transport, allowing LLM clients to instrument and diagnose running JVMs. + */ +public final class BTraceMcpServer { + private static final Logger log = LoggerFactory.getLogger(BTraceMcpServer.class); + private static final String SERVER_NAME = "btrace-mcp-server"; + private static final String SERVER_VERSION = "0.1.0"; + private static final String PROTOCOL_VERSION = "2024-11-05"; + + private final McpProtocol protocol; + + BTraceMcpServer(McpProtocol protocol) { + this.protocol = protocol; + } + + public static void main(String[] args) { + // Redirect System.out logging to stderr so stdout stays clean for MCP JSON-RPC + System.setProperty("org.slf4j.simpleLogger.logFile", "System.err"); + System.setProperty("org.slf4j.simpleLogger.defaultLogLevel", "warn"); + + log.info("Starting {}", SERVER_NAME); + McpProtocol proto = new McpProtocol(System.in, System.out); + BTraceMcpServer server = new BTraceMcpServer(proto); + try { + server.run(); + } catch (IOException e) { + log.error("Server error", e); + System.exit(1); + } + } + + @SuppressWarnings("unchecked") + void run() throws IOException { + while (true) { + Map message = protocol.readMessage(); + if (message == null) { + log.info("EOF on stdin, shutting down"); + break; + } + + String method = (String) message.get("method"); + Object id = message.get("id"); + Map params = + message.containsKey("params") ? (Map) message.get("params") : null; + + if (method == null) { + // Response or notification without method — ignore + continue; + } + + try { + switch (method) { + case "initialize": + handleInitialize(id); + break; + case "notifications/initialized": + // Client acknowledgement — no response needed + break; + case "tools/list": + handleToolsList(id); + break; + case "tools/call": + handleToolsCall(id, params); + break; + case "prompts/list": + handlePromptsList(id); + break; + case "prompts/get": + handlePromptsGet(id, params); + break; + default: + protocol.sendError(id, -32601, "Method not found: " + method); + } + } catch (Exception e) { + log.error("Error handling method: {}", method, e); + protocol.sendError(id, -32603, "Internal error: " + e.getMessage()); + } + } + } + + private void handleInitialize(Object id) throws IOException { + Map serverInfo = new LinkedHashMap<>(); + serverInfo.put("name", SERVER_NAME); + serverInfo.put("version", SERVER_VERSION); + + Map toolsCap = new LinkedHashMap<>(); + Map promptsCap = new LinkedHashMap<>(); + Map capabilities = new LinkedHashMap<>(); + capabilities.put("tools", toolsCap); + capabilities.put("prompts", promptsCap); + + Map result = new LinkedHashMap<>(); + result.put("protocolVersion", PROTOCOL_VERSION); + result.put("capabilities", capabilities); + result.put("serverInfo", serverInfo); + + protocol.sendResult(id, result); + } + + private void handleToolsList(Object id) throws IOException { + List tools = new ArrayList<>(); + tools.add(ListJvmsHandler.schema()); + tools.add(DeployOnelinerHandler.schema()); + tools.add(DeployScriptHandler.schema()); + tools.add(ListProbesHandler.schema()); + tools.add(SendEventHandler.schema()); + tools.add(DetachProbeHandler.schema()); + tools.add(ExitProbeHandler.schema()); + + Map result = new LinkedHashMap<>(); + result.put("tools", tools); + protocol.sendResult(id, result); + } + + @SuppressWarnings("unchecked") + private void handleToolsCall(Object id, Map params) throws IOException { + if (params == null) { + protocol.sendError(id, -32602, "Missing params"); + return; + } + String toolName = (String) params.get("name"); + Map arguments = + params.containsKey("arguments") + ? (Map) params.get("arguments") + : new LinkedHashMap<>(); + + Map result; + switch (toolName) { + case "list_jvms": + result = ListJvmsHandler.execute(arguments); + break; + case "deploy_oneliner": + result = DeployOnelinerHandler.execute(arguments); + break; + case "deploy_script": + result = DeployScriptHandler.execute(arguments); + break; + case "list_probes": + result = ListProbesHandler.execute(arguments); + break; + case "send_event": + result = SendEventHandler.execute(arguments); + break; + case "detach_probe": + result = DetachProbeHandler.execute(arguments); + break; + case "exit_probe": + result = ExitProbeHandler.execute(arguments); + break; + default: + protocol.sendError(id, -32602, "Unknown tool: " + toolName); + return; + } + protocol.sendResult(id, result); + } + + private void handlePromptsList(Object id) throws IOException { + Map result = new LinkedHashMap<>(); + result.put("prompts", DiagnosticPrompts.listPrompts()); + protocol.sendResult(id, result); + } + + @SuppressWarnings("unchecked") + private void handlePromptsGet(Object id, Map params) throws IOException { + if (params == null) { + protocol.sendError(id, -32602, "Missing params"); + return; + } + String name = (String) params.get("name"); + Map promptArgs = + params.containsKey("arguments") + ? (Map) params.get("arguments") + : new LinkedHashMap<>(); + + Map result = DiagnosticPrompts.getPrompt(name, promptArgs); + if (result == null) { + protocol.sendError(id, -32602, "Unknown prompt: " + name); + return; + } + protocol.sendResult(id, result); + } +} diff --git a/btrace-mcp-server/src/main/java/io/btrace/mcp/ClientManager.java b/btrace-mcp-server/src/main/java/io/btrace/mcp/ClientManager.java new file mode 100644 index 000000000..ce83cbcd5 --- /dev/null +++ b/btrace-mcp-server/src/main/java/io/btrace/mcp/ClientManager.java @@ -0,0 +1,77 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.mcp; + +import io.btrace.client.Client; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Manages BTrace client instances per JVM. Tracks active sessions so that tools like send_event, + * detach_probe, and exit_probe can interact with already-deployed probes. + */ +public final class ClientManager { + private static final Logger log = LoggerFactory.getLogger(ClientManager.class); + + /** Key is "pid:port", value is the active Client instance. */ + private static final Map activeClients = new ConcurrentHashMap<>(); + + private ClientManager() {} + + /** Creates a new BTrace Client for the given port. Does not reuse existing sessions. */ + public static Client getClient(int port) { + return new Client(port); + } + + /** Registers a client as active for a given PID and port. */ + public static void registerClient(String pid, int port, Client client) { + String key = pid + ":" + port; + activeClients.put(key, client); + log.debug("Registered client for {}", key); + } + + /** Returns an existing active client for the given PID and port, or null if none. */ + public static Client getExistingClient(String pid, int port) { + String key = pid + ":" + port; + return activeClients.get(key); + } + + /** Removes and returns an active client for the given PID and port. */ + public static Client removeClient(String pid, int port) { + String key = pid + ":" + port; + Client removed = activeClients.remove(key); + if (removed != null) { + log.debug("Removed client for {}", key); + } + return removed; + } + + /** Closes all active client sessions. Called on shutdown. */ + public static void closeAll() { + for (Map.Entry entry : activeClients.entrySet()) { + try { + entry.getValue().close(); + log.debug("Closed client for {}", entry.getKey()); + } catch (Exception e) { + log.warn("Error closing client for {}", entry.getKey(), e); + } + } + activeClients.clear(); + } +} diff --git a/btrace-mcp-server/src/main/java/io/btrace/mcp/McpProtocol.java b/btrace-mcp-server/src/main/java/io/btrace/mcp/McpProtocol.java new file mode 100644 index 000000000..a6b953f88 --- /dev/null +++ b/btrace-mcp-server/src/main/java/io/btrace/mcp/McpProtocol.java @@ -0,0 +1,435 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.mcp; + +import java.io.BufferedReader; +import java.io.IOException; +import java.io.InputStream; +import java.io.InputStreamReader; +import java.io.OutputStream; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Handles MCP JSON-RPC protocol over stdio. Reads JSON-RPC messages from stdin and writes responses + * to stdout. All logging goes to stderr. + */ +final class McpProtocol { + private static final Logger log = LoggerFactory.getLogger(McpProtocol.class); + private final BufferedReader reader; + private final OutputStream out; + + McpProtocol(InputStream in, OutputStream out) { + this.reader = new BufferedReader(new InputStreamReader(in, StandardCharsets.UTF_8)); + this.out = out; + } + + /** Reads a single JSON-RPC message from stdin. Returns null on EOF. */ + Map readMessage() throws IOException { + String line = reader.readLine(); + if (line == null) { + return null; + } + line = line.trim(); + if (line.isEmpty()) { + return null; + } + log.debug("Received: {}", line); + return parseJson(line); + } + + /** Writes a JSON-RPC response to stdout (one line, newline terminated). */ + synchronized void writeMessage(Map message) throws IOException { + String json = toJson(message); + log.debug("Sending: {}", json); + out.write(json.getBytes(StandardCharsets.UTF_8)); + out.write('\n'); + out.flush(); + } + + /** Sends a JSON-RPC success response. */ + void sendResult(Object id, Object result) throws IOException { + Map response = new LinkedHashMap<>(); + response.put("jsonrpc", "2.0"); + response.put("id", id); + response.put("result", result); + writeMessage(response); + } + + /** Sends a JSON-RPC error response. */ + void sendError(Object id, int code, String message) throws IOException { + sendError(id, code, message, null); + } + + /** Sends a JSON-RPC error response with optional data. */ + void sendError(Object id, int code, String message, Object data) throws IOException { + Map error = new LinkedHashMap<>(); + error.put("code", code); + error.put("message", message); + if (data != null) { + error.put("data", data); + } + Map response = new LinkedHashMap<>(); + response.put("jsonrpc", "2.0"); + response.put("id", id); + response.put("error", error); + writeMessage(response); + } + + // --- Minimal JSON parser (handles objects, arrays, strings, numbers, booleans, null) --- + + @SuppressWarnings("unchecked") + static Map parseJson(String json) { + Object result = new JsonParser(json.trim()).parseValue(); + if (result instanceof Map) { + return (Map) result; + } + throw new IllegalArgumentException("Expected JSON object, got: " + json); + } + + private static final class JsonParser { + private final String src; + private int pos; + + JsonParser(String src) { + this.src = src; + this.pos = 0; + } + + Object parseValue() { + skipWhitespace(); + if (pos >= src.length()) { + throw new IllegalArgumentException("Unexpected end of JSON"); + } + char c = src.charAt(pos); + if (c == '{') { + return parseObject(); + } + if (c == '[') { + return parseArray(); + } + if (c == '"') { + return parseString(); + } + if (c == 't' || c == 'f') { + return parseBoolean(); + } + if (c == 'n') { + return parseNull(); + } + return parseNumber(); + } + + Map parseObject() { + expect('{'); + Map map = new LinkedHashMap<>(); + skipWhitespace(); + if (pos < src.length() && src.charAt(pos) == '}') { + pos++; + return map; + } + while (true) { + skipWhitespace(); + String key = parseString(); + skipWhitespace(); + expect(':'); + Object value = parseValue(); + map.put(key, value); + skipWhitespace(); + if (pos < src.length() && src.charAt(pos) == ',') { + pos++; + } else { + break; + } + } + expect('}'); + return map; + } + + List parseArray() { + expect('['); + List list = new ArrayList<>(); + skipWhitespace(); + if (pos < src.length() && src.charAt(pos) == ']') { + pos++; + return list; + } + while (true) { + list.add(parseValue()); + skipWhitespace(); + if (pos < src.length() && src.charAt(pos) == ',') { + pos++; + } else { + break; + } + } + expect(']'); + return list; + } + + String parseString() { + expect('"'); + StringBuilder sb = new StringBuilder(); + while (pos < src.length()) { + char c = src.charAt(pos++); + if (c == '"') { + return sb.toString(); + } + if (c == '\\') { + if (pos >= src.length()) { + break; + } + char esc = src.charAt(pos++); + switch (esc) { + case '"': + case '\\': + case '/': + sb.append(esc); + break; + case 'b': + sb.append('\b'); + break; + case 'f': + sb.append('\f'); + break; + case 'n': + sb.append('\n'); + break; + case 'r': + sb.append('\r'); + break; + case 't': + sb.append('\t'); + break; + case 'u': + if (pos + 4 <= src.length()) { + String hex = src.substring(pos, pos + 4); + sb.append((char) Integer.parseInt(hex, 16)); + pos += 4; + } + break; + default: + sb.append(esc); + } + } else { + sb.append(c); + } + } + throw new IllegalArgumentException("Unterminated string"); + } + + Object parseNumber() { + int start = pos; + if (pos < src.length() && src.charAt(pos) == '-') { + pos++; + } + while (pos < src.length() && Character.isDigit(src.charAt(pos))) { + pos++; + } + boolean isFloat = false; + if (pos < src.length() && src.charAt(pos) == '.') { + isFloat = true; + pos++; + while (pos < src.length() && Character.isDigit(src.charAt(pos))) { + pos++; + } + } + if (pos < src.length() && (src.charAt(pos) == 'e' || src.charAt(pos) == 'E')) { + isFloat = true; + pos++; + if (pos < src.length() && (src.charAt(pos) == '+' || src.charAt(pos) == '-')) { + pos++; + } + while (pos < src.length() && Character.isDigit(src.charAt(pos))) { + pos++; + } + } + String numStr = src.substring(start, pos); + if (isFloat) { + try { + return Double.parseDouble(numStr); + } catch (NumberFormatException e) { + throw new IllegalArgumentException("Invalid number at " + start + ": " + numStr, e); + } + } + long val; + try { + val = Long.parseLong(numStr); + } catch (NumberFormatException e) { + throw new IllegalArgumentException("Invalid number at " + start + ": " + numStr, e); + } + if (val >= Integer.MIN_VALUE && val <= Integer.MAX_VALUE) { + return (int) val; + } + return val; + } + + Object parseBoolean() { + if (src.startsWith("true", pos)) { + pos += 4; + return Boolean.TRUE; + } + if (src.startsWith("false", pos)) { + pos += 5; + return Boolean.FALSE; + } + throw new IllegalArgumentException("Invalid boolean at " + pos); + } + + Object parseNull() { + if (src.startsWith("null", pos)) { + pos += 4; + return null; + } + throw new IllegalArgumentException("Invalid null at " + pos); + } + + void skipWhitespace() { + while (pos < src.length() && Character.isWhitespace(src.charAt(pos))) { + pos++; + } + } + + void expect(char c) { + skipWhitespace(); + if (pos >= src.length() || src.charAt(pos) != c) { + throw new IllegalArgumentException( + "Expected '" + + c + + "' at " + + pos + + " but got: " + + (pos < src.length() ? src.charAt(pos) : "EOF")); + } + pos++; + } + } + + // --- Minimal JSON serializer --- + + @SuppressWarnings("unchecked") + static String toJson(Object obj) { + if (obj == null) { + return "null"; + } + if (obj instanceof String) { + return escapeJsonString((String) obj); + } + if (obj instanceof Number || obj instanceof Boolean) { + return obj.toString(); + } + if (obj instanceof Map) { + Map map = (Map) obj; + StringBuilder sb = new StringBuilder("{"); + boolean first = true; + for (Map.Entry entry : map.entrySet()) { + if (!first) { + sb.append(","); + } + sb.append(escapeJsonString(entry.getKey())); + sb.append(":"); + sb.append(toJson(entry.getValue())); + first = false; + } + sb.append("}"); + return sb.toString(); + } + if (obj instanceof List) { + List list = (List) obj; + StringBuilder sb = new StringBuilder("["); + for (int i = 0; i < list.size(); i++) { + if (i > 0) { + sb.append(","); + } + sb.append(toJson(list.get(i))); + } + sb.append("]"); + return sb.toString(); + } + if (obj instanceof Object[]) { + Object[] arr = (Object[]) obj; + StringBuilder sb = new StringBuilder("["); + for (int i = 0; i < arr.length; i++) { + if (i > 0) { + sb.append(","); + } + sb.append(toJson(arr[i])); + } + sb.append("]"); + return sb.toString(); + } + return escapeJsonString(obj.toString()); + } + + private static String escapeJsonString(String s) { + StringBuilder sb = new StringBuilder("\""); + for (int i = 0; i < s.length(); i++) { + char c = s.charAt(i); + switch (c) { + case '"': + sb.append("\\\""); + break; + case '\\': + sb.append("\\\\"); + break; + case '\b': + sb.append("\\b"); + break; + case '\f': + sb.append("\\f"); + break; + case '\n': + sb.append("\\n"); + break; + case '\r': + sb.append("\\r"); + break; + case '\t': + sb.append("\\t"); + break; + default: + if (c < 0x20) { + sb.append(String.format("\\u%04x", (int) c)); + } else { + sb.append(c); + } + } + } + sb.append("\""); + return sb.toString(); + } + + /** Helper to build a tool content result (text). */ + static Map toolResult(String text, boolean isError) { + Map content = new LinkedHashMap<>(); + content.put("type", "text"); + content.put("text", text); + + List contentList = new ArrayList<>(); + contentList.add(content); + + Map result = new LinkedHashMap<>(); + result.put("content", contentList); + result.put("isError", isError); + return result; + } +} diff --git a/btrace-mcp-server/src/main/java/io/btrace/mcp/prompts/DiagnosticPrompts.java b/btrace-mcp-server/src/main/java/io/btrace/mcp/prompts/DiagnosticPrompts.java new file mode 100644 index 000000000..38efdf60e --- /dev/null +++ b/btrace-mcp-server/src/main/java/io/btrace/mcp/prompts/DiagnosticPrompts.java @@ -0,0 +1,312 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.mcp.prompts; + +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +/** Provides MCP prompt templates for common BTrace diagnostic scenarios. */ +public final class DiagnosticPrompts { + + private DiagnosticPrompts() {} + + /** Returns all available prompt schemas for prompts/list. */ + public static List> listPrompts() { + List> prompts = new ArrayList<>(); + prompts.add(diagnoseSlowEndpointSchema()); + prompts.add(findExceptionSourceSchema()); + prompts.add(profileMethodSchema()); + return prompts; + } + + /** Returns a prompt by name, or null if not found. */ + public static Map getPrompt(String name, Map arguments) { + switch (name) { + case "diagnose_slow_endpoint": + return diagnoseSlowEndpoint(arguments); + case "find_exception_source": + return findExceptionSource(arguments); + case "profile_method": + return profileMethod(arguments); + default: + return null; + } + } + + // --- diagnose_slow_endpoint --- + + private static Map diagnoseSlowEndpointSchema() { + Map prompt = new LinkedHashMap<>(); + prompt.put("name", "diagnose_slow_endpoint"); + prompt.put( + "description", + "Step-by-step guide to diagnose a slow HTTP endpoint using BTrace. " + + "Helps identify where time is spent in request processing."); + + List> args = new ArrayList<>(); + args.add( + promptArg("endpoint_class", "Fully qualified class name of the endpoint/controller", true)); + args.add(promptArg("endpoint_method", "Method name to diagnose", true)); + args.add( + promptArg( + "pid", "PID of the target JVM (optional, will use list_jvms if not provided)", false)); + prompt.put("arguments", args); + return prompt; + } + + private static Map diagnoseSlowEndpoint(Map arguments) { + String endpointClass = getArg(arguments, "endpoint_class", "com.example.Controller"); + String endpointMethod = getArg(arguments, "endpoint_method", "handleRequest"); + String pid = getArg(arguments, "pid", null); + + StringBuilder text = new StringBuilder(); + text.append("# Diagnosing Slow Endpoint: ") + .append(endpointClass) + .append("::") + .append(endpointMethod) + .append("\n\n"); + + text.append("## Step 1: Find the target JVM\n"); + if (pid != null) { + text.append("Target PID: ").append(pid).append("\n"); + } else { + text.append("Use the `list_jvms` tool to find the target JVM PID.\n"); + } + text.append("\n"); + + text.append("## Step 2: Measure endpoint latency\n"); + text.append("Deploy a oneliner to measure the method's execution time:\n"); + text.append("```\n"); + text.append(endpointClass) + .append("::") + .append(endpointMethod) + .append(" @return { print method, duration }\n"); + text.append("```\n\n"); + + text.append("## Step 3: Trace internal method calls\n"); + text.append( + "If the endpoint is slow, trace the internal methods it calls to find the bottleneck:\n"); + text.append("```\n"); + text.append(endpointClass) + .append("::/.*/") + .append(" @return if duration>10ms { print method, duration }\n"); + text.append("```\n\n"); + + text.append("## Step 4: Analyze results\n"); + text.append("Look for methods with unexpectedly high durations. Common causes include:\n"); + text.append("- Database queries taking too long\n"); + text.append("- External service calls with high latency\n"); + text.append("- Lock contention\n"); + text.append("- Excessive object allocation\n\n"); + + text.append("## Step 5: Clean up\n"); + text.append("Use `exit_probe` to remove the instrumentation when done.\n"); + + return promptResult(text.toString()); + } + + // --- find_exception_source --- + + private static Map findExceptionSourceSchema() { + Map prompt = new LinkedHashMap<>(); + prompt.put("name", "find_exception_source"); + prompt.put( + "description", + "Guide to find where specific exceptions originate in a running JVM. " + + "Uses BTrace to intercept exception constructors and capture stack traces."); + + List> args = new ArrayList<>(); + args.add( + promptArg("exception_class", "Exception class name (e.g. NullPointerException)", true)); + args.add(promptArg("pid", "PID of the target JVM (optional)", false)); + prompt.put("arguments", args); + return prompt; + } + + private static Map findExceptionSource(Map arguments) { + String exceptionClass = getArg(arguments, "exception_class", "java.lang.NullPointerException"); + String pid = getArg(arguments, "pid", null); + + StringBuilder text = new StringBuilder(); + text.append("# Finding Exception Source: ").append(exceptionClass).append("\n\n"); + + text.append("## Step 1: Find the target JVM\n"); + if (pid != null) { + text.append("Target PID: ").append(pid).append("\n"); + } else { + text.append("Use the `list_jvms` tool to find the target JVM PID.\n"); + } + text.append("\n"); + + text.append("## Step 2: Deploy an exception tracing script\n"); + text.append("Deploy this BTrace script to capture exception creation with stack traces:\n\n"); + text.append("```java\n"); + text.append("import io.btrace.core.annotations.*;\n"); + text.append("import static io.btrace.core.BTraceUtils.*;\n\n"); + text.append("@BTrace\n"); + text.append("public class ExceptionTracer {\n"); + text.append(" @OnMethod(\n"); + text.append(" clazz = \"").append(exceptionClass).append("\",\n"); + text.append(" method = \"\"\n"); + text.append(" )\n"); + text.append(" public static void onException(@Self Throwable self) {\n"); + text.append(" println(\"--- Exception created: \" + Strings.str(self) + \" ---\");\n"); + text.append(" Threads.jstack();\n"); + text.append(" println(\"\");\n"); + text.append(" }\n"); + text.append("}\n"); + text.append("```\n\n"); + + text.append("## Step 3: Analyze the stack traces\n"); + text.append( + "Each time the exception is created, you will see the full stack trace " + + "showing exactly which code path creates it.\n\n"); + + text.append("## Step 4: Clean up\n"); + text.append("Use `exit_probe` to remove the instrumentation when done.\n"); + + return promptResult(text.toString()); + } + + // --- profile_method --- + + private static Map profileMethodSchema() { + Map prompt = new LinkedHashMap<>(); + prompt.put("name", "profile_method"); + prompt.put( + "description", + "Guide to profile a specific method's latency distribution using BTrace. " + + "Captures timing data to understand performance characteristics."); + + List> args = new ArrayList<>(); + args.add(promptArg("class_name", "Fully qualified class name", true)); + args.add(promptArg("method_name", "Method name to profile", true)); + args.add(promptArg("pid", "PID of the target JVM (optional)", false)); + prompt.put("arguments", args); + return prompt; + } + + private static Map profileMethod(Map arguments) { + String className = getArg(arguments, "class_name", "com.example.Service"); + String methodName = getArg(arguments, "method_name", "process"); + String pid = getArg(arguments, "pid", null); + + StringBuilder text = new StringBuilder(); + text.append("# Profiling Method: ") + .append(className) + .append("::") + .append(methodName) + .append("\n\n"); + + text.append("## Step 1: Find the target JVM\n"); + if (pid != null) { + text.append("Target PID: ").append(pid).append("\n"); + } else { + text.append("Use the `list_jvms` tool to find the target JVM PID.\n"); + } + text.append("\n"); + + text.append("## Step 2: Quick latency check with oneliner\n"); + text.append("Start with a simple oneliner to see individual call durations:\n"); + text.append("```\n"); + text.append(className) + .append("::") + .append(methodName) + .append(" @return { print method, duration }\n"); + text.append("```\n\n"); + + text.append("## Step 3: Detailed profiling with histogram\n"); + text.append("For a latency distribution, deploy this BTrace script:\n\n"); + text.append("```java\n"); + text.append("import io.btrace.core.annotations.*;\n"); + text.append("import io.btrace.core.BTraceUtils;\n"); + text.append("import static io.btrace.core.BTraceUtils.*;\n\n"); + text.append("@BTrace\n"); + text.append("public class MethodProfiler {\n"); + text.append(" private static long count;\n"); + text.append(" private static long totalTime;\n\n"); + text.append(" @OnMethod(\n"); + text.append(" clazz = \"").append(className).append("\",\n"); + text.append(" method = \"").append(methodName).append("\",\n"); + text.append(" location = @Location(Kind.RETURN)\n"); + text.append(" )\n"); + text.append(" public static void onReturn(@Duration long duration) {\n"); + text.append(" count++;\n"); + text.append(" totalTime += duration;\n"); + text.append(" println(\"Call #\" + count + \": \" + (duration / 1000000) + \"ms\");\n"); + text.append(" }\n\n"); + text.append(" @OnEvent\n"); + text.append(" public static void onEvent() {\n"); + text.append(" println(\"=== Summary ===\");\n"); + text.append(" println(\"Total calls: \" + count);\n"); + text.append(" if (count > 0) {\n"); + text.append( + " println(\"Avg duration: \" + ((totalTime / count) / 1000000) + \"ms\");\n"); + text.append(" }\n"); + text.append(" }\n"); + text.append("}\n"); + text.append("```\n\n"); + + text.append("## Step 4: Get summary\n"); + text.append("Use `send_event` to trigger the @OnEvent handler and get a summary.\n\n"); + + text.append("## Step 5: Clean up\n"); + text.append("Use `exit_probe` to remove the instrumentation when done.\n"); + + return promptResult(text.toString()); + } + + // --- Helpers --- + + private static Map promptArg(String name, String description, boolean required) { + Map arg = new LinkedHashMap<>(); + arg.put("name", name); + arg.put("description", description); + arg.put("required", required); + return arg; + } + + private static String getArg(Map args, String key, String defaultVal) { + if (args == null) { + return defaultVal; + } + Object val = args.get(key); + if (val == null) { + return defaultVal; + } + return val.toString(); + } + + private static Map promptResult(String text) { + Map message = new LinkedHashMap<>(); + message.put("role", "user"); + + Map content = new LinkedHashMap<>(); + content.put("type", "text"); + content.put("text", text); + message.put("content", content); + + List messages = new ArrayList<>(); + messages.add(message); + + Map result = new LinkedHashMap<>(); + result.put("messages", messages); + return result; + } +} diff --git a/btrace-mcp-server/src/main/java/io/btrace/mcp/tools/DeployOnelinerHandler.java b/btrace-mcp-server/src/main/java/io/btrace/mcp/tools/DeployOnelinerHandler.java new file mode 100644 index 000000000..f6ee10121 --- /dev/null +++ b/btrace-mcp-server/src/main/java/io/btrace/mcp/tools/DeployOnelinerHandler.java @@ -0,0 +1,223 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.mcp.tools; + +import io.btrace.client.Client; +import io.btrace.compiler.oneliner.OnelinerAST.OnelinerNode; +import io.btrace.compiler.oneliner.OnelinerCodeGenerator; +import io.btrace.compiler.oneliner.OnelinerParser; +import io.btrace.compiler.oneliner.OnelinerValidator; +import io.btrace.core.comm.Command; +import io.btrace.core.comm.PrintableCommand; +import io.btrace.mcp.ClientManager; +import java.io.PrintWriter; +import java.io.StringWriter; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Handles the deploy_oneliner MCP tool - deploys a BTrace oneliner probe. */ +public final class DeployOnelinerHandler { + private static final Logger log = LoggerFactory.getLogger(DeployOnelinerHandler.class); + private static final int DEFAULT_PORT = 2020; + private static final int PROBE_TIMEOUT_SECONDS = 30; + + private DeployOnelinerHandler() {} + + /** Returns tool schema for MCP tools/list. */ + public static Map schema() { + Map tool = new LinkedHashMap<>(); + tool.put("name", "deploy_oneliner"); + tool.put( + "description", + "Deploy a BTrace oneliner probe to a running JVM. " + + "Oneliners are concise probe expressions like: " + + "\"com.example.Service::method @return { print duration }\" " + + "or \"com.example.Dao::query @return if duration>100ms { print method, duration }\". " + + "The probe attaches to the target JVM and captures output for the specified duration."); + + Map properties = new LinkedHashMap<>(); + + Map pidProp = new LinkedHashMap<>(); + pidProp.put("type", "string"); + pidProp.put("description", "PID of the target JVM (use list_jvms to find it)"); + properties.put("pid", pidProp); + + Map onelinerProp = new LinkedHashMap<>(); + onelinerProp.put("type", "string"); + onelinerProp.put( + "description", + "BTrace oneliner expression, e.g. " + + "\"com.example.Service::method @return { print method, duration }\""); + properties.put("oneliner", onelinerProp); + + Map portProp = new LinkedHashMap<>(); + portProp.put("type", "integer"); + portProp.put("description", "BTrace agent port (default: 2020)"); + properties.put("port", portProp); + + List required = new ArrayList<>(); + required.add("pid"); + required.add("oneliner"); + + Map inputSchema = new LinkedHashMap<>(); + inputSchema.put("type", "object"); + inputSchema.put("properties", properties); + inputSchema.put("required", required); + tool.put("inputSchema", inputSchema); + return tool; + } + + /** Executes the deploy_oneliner tool. */ + @SuppressWarnings("unchecked") + public static Map execute(Map arguments) { + String pid = getStringArg(arguments, "pid"); + String oneliner = getStringArg(arguments, "oneliner"); + int port = getIntArg(arguments, "port", DEFAULT_PORT); + + if (pid == null || pid.isEmpty()) { + return toolResult("Error: 'pid' parameter is required", true); + } + if (oneliner == null || oneliner.isEmpty()) { + return toolResult("Error: 'oneliner' parameter is required", true); + } + + try { + // Parse and validate the oneliner + OnelinerNode ast = OnelinerParser.parse(oneliner); + OnelinerValidator.validate(ast, oneliner); + String className = "BTraceOneliner_" + System.currentTimeMillis(); + String javaSource = OnelinerCodeGenerator.generate(ast, className); + String fileName = className + ".java"; + + log.info("Generated oneliner source for {}: {}", oneliner, javaSource); + + // Compile the oneliner + Client client = ClientManager.getClient(port); + StringWriter errorWriter = new StringWriter(); + PrintWriter errPw = new PrintWriter(errorWriter); + byte[] code = client.compileSource(fileName, javaSource, ".", errPw, null); + + if (code == null) { + String errors = errorWriter.toString(); + return toolResult( + "Oneliner compilation failed:\n" + errors + "\nGenerated source:\n" + javaSource, true); + } + + // Attach and submit + client.attach(pid, null, "."); + + StringBuilder output = new StringBuilder(); + CountDownLatch statusLatch = new CountDownLatch(1); + AtomicBoolean success = new AtomicBoolean(false); + AtomicBoolean exited = new AtomicBoolean(false); + + client.submit( + "localhost", + fileName, + code, + new String[0], + cmd -> { + int type = cmd.getType(); + if (cmd instanceof PrintableCommand) { + StringWriter sw = new StringWriter(); + ((PrintableCommand) cmd).print(new java.io.PrintWriter(sw)); + output.append(sw.toString()); + } + if (type == Command.STATUS) { + success.set(true); + statusLatch.countDown(); + } + if (type == Command.EXIT) { + exited.set(true); + statusLatch.countDown(); + } + }); + + // Wait for probe to start (or fail) + boolean started = statusLatch.await(PROBE_TIMEOUT_SECONDS, TimeUnit.SECONDS); + + if (!started) { + return toolResult("Probe deployment timed out after " + PROBE_TIMEOUT_SECONDS + "s", true); + } + + if (exited.get() && !success.get()) { + return toolResult("Probe exited with error:\n" + output.toString(), true); + } + + ClientManager.registerClient(pid, port, client); + + String resultText = + "Probe deployed successfully to PID " + + pid + + ".\n" + + "Oneliner: " + + oneliner + + "\n" + + "Port: " + + port + + "\n"; + if (output.length() > 0) { + resultText += "\nInitial output:\n" + output.toString(); + } + resultText += + "\nUse send_event, detach_probe, or exit_probe to interact with the running probe."; + return toolResult(resultText, false); + } catch (Exception e) { + log.error("Failed to deploy oneliner", e); + return toolResult("Error deploying oneliner: " + e.getMessage(), true); + } + } + + private static String getStringArg(Map args, String key) { + Object val = args == null ? null : args.get(key); + return val == null ? null : val.toString(); + } + + private static int getIntArg(Map args, String key, int defaultVal) { + Object val = args == null ? null : args.get(key); + if (val == null) { + return defaultVal; + } + if (val instanceof Number) { + return ((Number) val).intValue(); + } + try { + return Integer.parseInt(val.toString()); + } catch (NumberFormatException e) { + return defaultVal; + } + } + + private static Map toolResult(String text, boolean isError) { + Map content = new LinkedHashMap<>(); + content.put("type", "text"); + content.put("text", text); + List contentList = new ArrayList<>(); + contentList.add(content); + Map result = new LinkedHashMap<>(); + result.put("content", contentList); + result.put("isError", isError); + return result; + } +} diff --git a/btrace-mcp-server/src/main/java/io/btrace/mcp/tools/DeployScriptHandler.java b/btrace-mcp-server/src/main/java/io/btrace/mcp/tools/DeployScriptHandler.java new file mode 100644 index 000000000..e19296f86 --- /dev/null +++ b/btrace-mcp-server/src/main/java/io/btrace/mcp/tools/DeployScriptHandler.java @@ -0,0 +1,218 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.mcp.tools; + +import io.btrace.client.Client; +import io.btrace.core.comm.Command; +import io.btrace.core.comm.PrintableCommand; +import io.btrace.mcp.ClientManager; +import java.io.PrintWriter; +import java.io.StringWriter; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Handles the deploy_script MCP tool - deploys a full BTrace Java script. */ +public final class DeployScriptHandler { + private static final Logger log = LoggerFactory.getLogger(DeployScriptHandler.class); + private static final int DEFAULT_PORT = 2020; + private static final int PROBE_TIMEOUT_SECONDS = 30; + + private DeployScriptHandler() {} + + /** Returns tool schema for MCP tools/list. */ + public static Map schema() { + Map tool = new LinkedHashMap<>(); + tool.put("name", "deploy_script"); + tool.put( + "description", + "Deploy a full BTrace Java script to a running JVM. " + + "The script must be a valid BTrace program with @BTrace annotation. " + + "Use this for complex instrumentation that cannot be expressed as a oneliner."); + + Map properties = new LinkedHashMap<>(); + + Map pidProp = new LinkedHashMap<>(); + pidProp.put("type", "string"); + pidProp.put("description", "PID of the target JVM (use list_jvms to find it)"); + properties.put("pid", pidProp); + + Map scriptProp = new LinkedHashMap<>(); + scriptProp.put("type", "string"); + scriptProp.put( + "description", + "Full BTrace Java source code. Must include @BTrace annotation and proper imports."); + properties.put("script", scriptProp); + + Map argsProp = new LinkedHashMap<>(); + argsProp.put("type", "array"); + Map argsItems = new LinkedHashMap<>(); + argsItems.put("type", "string"); + argsProp.put("items", argsItems); + argsProp.put("description", "Optional arguments to pass to the BTrace script"); + properties.put("args", argsProp); + + Map portProp = new LinkedHashMap<>(); + portProp.put("type", "integer"); + portProp.put("description", "BTrace agent port (default: 2020)"); + properties.put("port", portProp); + + List required = new ArrayList<>(); + required.add("pid"); + required.add("script"); + + Map inputSchema = new LinkedHashMap<>(); + inputSchema.put("type", "object"); + inputSchema.put("properties", properties); + inputSchema.put("required", required); + tool.put("inputSchema", inputSchema); + return tool; + } + + /** Executes the deploy_script tool. */ + @SuppressWarnings("unchecked") + public static Map execute(Map arguments) { + String pid = getStringArg(arguments, "pid"); + String script = getStringArg(arguments, "script"); + int port = getIntArg(arguments, "port", DEFAULT_PORT); + + if (pid == null || pid.isEmpty()) { + return toolResult("Error: 'pid' parameter is required", true); + } + if (script == null || script.isEmpty()) { + return toolResult("Error: 'script' parameter is required", true); + } + + // Extract args + String[] btraceArgs = new String[0]; + Object argsObj = arguments == null ? null : arguments.get("args"); + if (argsObj instanceof List) { + List argsList = (List) argsObj; + btraceArgs = new String[argsList.size()]; + for (int i = 0; i < argsList.size(); i++) { + btraceArgs[i] = argsList.get(i).toString(); + } + } + + try { + String fileName = "BTraceScript_" + System.currentTimeMillis() + ".java"; + + // Compile the script + Client client = ClientManager.getClient(port); + StringWriter errorWriter = new StringWriter(); + PrintWriter errPw = new PrintWriter(errorWriter); + byte[] code = client.compileSource(fileName, script, ".", errPw, null); + + if (code == null) { + String errors = errorWriter.toString(); + return toolResult("Script compilation failed:\n" + errors, true); + } + + // Attach and submit + client.attach(pid, null, "."); + + StringBuilder output = new StringBuilder(); + CountDownLatch statusLatch = new CountDownLatch(1); + AtomicBoolean success = new AtomicBoolean(false); + AtomicBoolean exited = new AtomicBoolean(false); + final String[] finalArgs = btraceArgs; + + client.submit( + "localhost", + fileName, + code, + finalArgs, + cmd -> { + int type = cmd.getType(); + if (cmd instanceof PrintableCommand) { + StringWriter sw = new StringWriter(); + ((PrintableCommand) cmd).print(new java.io.PrintWriter(sw)); + output.append(sw.toString()); + } + if (type == Command.STATUS) { + success.set(true); + statusLatch.countDown(); + } + if (type == Command.EXIT) { + exited.set(true); + statusLatch.countDown(); + } + }); + + boolean started = statusLatch.await(PROBE_TIMEOUT_SECONDS, TimeUnit.SECONDS); + + if (!started) { + return toolResult("Probe deployment timed out after " + PROBE_TIMEOUT_SECONDS + "s", true); + } + + if (exited.get() && !success.get()) { + return toolResult("Probe exited with error:\n" + output.toString(), true); + } + + ClientManager.registerClient(pid, port, client); + + String resultText = "Script deployed successfully to PID " + pid + ".\nPort: " + port + "\n"; + if (output.length() > 0) { + resultText += "\nInitial output:\n" + output.toString(); + } + resultText += + "\nUse send_event, detach_probe, or exit_probe to interact with the running probe."; + return toolResult(resultText, false); + } catch (Exception e) { + log.error("Failed to deploy script", e); + return toolResult("Error deploying script: " + e.getMessage(), true); + } + } + + private static String getStringArg(Map args, String key) { + Object val = args == null ? null : args.get(key); + return val == null ? null : val.toString(); + } + + private static int getIntArg(Map args, String key, int defaultVal) { + Object val = args == null ? null : args.get(key); + if (val == null) { + return defaultVal; + } + if (val instanceof Number) { + return ((Number) val).intValue(); + } + try { + return Integer.parseInt(val.toString()); + } catch (NumberFormatException e) { + return defaultVal; + } + } + + private static Map toolResult(String text, boolean isError) { + Map content = new LinkedHashMap<>(); + content.put("type", "text"); + content.put("text", text); + List contentList = new ArrayList<>(); + contentList.add(content); + Map result = new LinkedHashMap<>(); + result.put("content", contentList); + result.put("isError", isError); + return result; + } +} diff --git a/btrace-mcp-server/src/main/java/io/btrace/mcp/tools/DetachProbeHandler.java b/btrace-mcp-server/src/main/java/io/btrace/mcp/tools/DetachProbeHandler.java new file mode 100644 index 000000000..ddcf45f23 --- /dev/null +++ b/btrace-mcp-server/src/main/java/io/btrace/mcp/tools/DetachProbeHandler.java @@ -0,0 +1,125 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.mcp.tools; + +import io.btrace.client.Client; +import io.btrace.mcp.ClientManager; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Handles the detach_probe MCP tool - detaches the client from a running probe (probe continues + * running). + */ +public final class DetachProbeHandler { + private static final Logger log = LoggerFactory.getLogger(DetachProbeHandler.class); + private static final int DEFAULT_PORT = 2020; + + private DetachProbeHandler() {} + + /** Returns tool schema for MCP tools/list. */ + public static Map schema() { + Map tool = new LinkedHashMap<>(); + tool.put("name", "detach_probe"); + tool.put( + "description", + "Detach from a running BTrace probe. The probe continues running on the target JVM " + + "but the client disconnects. You can later reconnect to it."); + + Map properties = new LinkedHashMap<>(); + + Map pidProp = new LinkedHashMap<>(); + pidProp.put("type", "string"); + pidProp.put("description", "PID of the target JVM"); + properties.put("pid", pidProp); + + Map portProp = new LinkedHashMap<>(); + portProp.put("type", "integer"); + portProp.put("description", "BTrace agent port (default: 2020)"); + properties.put("port", portProp); + + List required = new ArrayList<>(); + required.add("pid"); + + Map inputSchema = new LinkedHashMap<>(); + inputSchema.put("type", "object"); + inputSchema.put("properties", properties); + inputSchema.put("required", required); + tool.put("inputSchema", inputSchema); + return tool; + } + + /** Executes the detach_probe tool. */ + public static Map execute(Map arguments) { + String pid = getStringArg(arguments, "pid"); + int port = getIntArg(arguments, "port", DEFAULT_PORT); + + if (pid == null || pid.isEmpty()) { + return toolResult("Error: 'pid' parameter is required", true); + } + + try { + Client client = ClientManager.getExistingClient(pid, port); + if (client == null) { + return toolResult("No active BTrace session for PID " + pid, true); + } + + client.sendDisconnect(); + ClientManager.removeClient(pid, port); + return toolResult("Detached from probe on PID " + pid + ". Probe continues running.", false); + } catch (Exception e) { + log.error("Failed to detach probe", e); + return toolResult("Error detaching from PID " + pid + ": " + e.getMessage(), true); + } + } + + private static String getStringArg(Map args, String key) { + Object val = args == null ? null : args.get(key); + return val == null ? null : val.toString(); + } + + private static int getIntArg(Map args, String key, int defaultVal) { + Object val = args == null ? null : args.get(key); + if (val == null) { + return defaultVal; + } + if (val instanceof Number) { + return ((Number) val).intValue(); + } + try { + return Integer.parseInt(val.toString()); + } catch (NumberFormatException e) { + return defaultVal; + } + } + + private static Map toolResult(String text, boolean isError) { + Map content = new LinkedHashMap<>(); + content.put("type", "text"); + content.put("text", text); + List contentList = new ArrayList<>(); + contentList.add(content); + Map result = new LinkedHashMap<>(); + result.put("content", contentList); + result.put("isError", isError); + return result; + } +} diff --git a/btrace-mcp-server/src/main/java/io/btrace/mcp/tools/ExitProbeHandler.java b/btrace-mcp-server/src/main/java/io/btrace/mcp/tools/ExitProbeHandler.java new file mode 100644 index 000000000..362cae084 --- /dev/null +++ b/btrace-mcp-server/src/main/java/io/btrace/mcp/tools/ExitProbeHandler.java @@ -0,0 +1,122 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.mcp.tools; + +import io.btrace.client.Client; +import io.btrace.mcp.ClientManager; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Handles the exit_probe MCP tool - stops and removes a probe from the target JVM. */ +public final class ExitProbeHandler { + private static final Logger log = LoggerFactory.getLogger(ExitProbeHandler.class); + private static final int DEFAULT_PORT = 2020; + + private ExitProbeHandler() {} + + /** Returns tool schema for MCP tools/list. */ + public static Map schema() { + Map tool = new LinkedHashMap<>(); + tool.put("name", "exit_probe"); + tool.put( + "description", + "Stop and remove a BTrace probe from the target JVM. " + + "This terminates the probe and removes all instrumentation."); + + Map properties = new LinkedHashMap<>(); + + Map pidProp = new LinkedHashMap<>(); + pidProp.put("type", "string"); + pidProp.put("description", "PID of the target JVM"); + properties.put("pid", pidProp); + + Map portProp = new LinkedHashMap<>(); + portProp.put("type", "integer"); + portProp.put("description", "BTrace agent port (default: 2020)"); + properties.put("port", portProp); + + List required = new ArrayList<>(); + required.add("pid"); + + Map inputSchema = new LinkedHashMap<>(); + inputSchema.put("type", "object"); + inputSchema.put("properties", properties); + inputSchema.put("required", required); + tool.put("inputSchema", inputSchema); + return tool; + } + + /** Executes the exit_probe tool. */ + public static Map execute(Map arguments) { + String pid = getStringArg(arguments, "pid"); + int port = getIntArg(arguments, "port", DEFAULT_PORT); + + if (pid == null || pid.isEmpty()) { + return toolResult("Error: 'pid' parameter is required", true); + } + + try { + Client client = ClientManager.getExistingClient(pid, port); + if (client == null) { + return toolResult("No active BTrace session for PID " + pid, true); + } + + client.sendExit(0); + ClientManager.removeClient(pid, port); + return toolResult("Probe stopped and removed from PID " + pid, false); + } catch (Exception e) { + log.error("Failed to exit probe", e); + return toolResult("Error stopping probe on PID " + pid + ": " + e.getMessage(), true); + } + } + + private static String getStringArg(Map args, String key) { + Object val = args == null ? null : args.get(key); + return val == null ? null : val.toString(); + } + + private static int getIntArg(Map args, String key, int defaultVal) { + Object val = args == null ? null : args.get(key); + if (val == null) { + return defaultVal; + } + if (val instanceof Number) { + return ((Number) val).intValue(); + } + try { + return Integer.parseInt(val.toString()); + } catch (NumberFormatException e) { + return defaultVal; + } + } + + private static Map toolResult(String text, boolean isError) { + Map content = new LinkedHashMap<>(); + content.put("type", "text"); + content.put("text", text); + List contentList = new ArrayList<>(); + contentList.add(content); + Map result = new LinkedHashMap<>(); + result.put("content", contentList); + result.put("isError", isError); + return result; + } +} diff --git a/btrace-mcp-server/src/main/java/io/btrace/mcp/tools/ListJvmsHandler.java b/btrace-mcp-server/src/main/java/io/btrace/mcp/tools/ListJvmsHandler.java new file mode 100644 index 000000000..10e5b3898 --- /dev/null +++ b/btrace-mcp-server/src/main/java/io/btrace/mcp/tools/ListJvmsHandler.java @@ -0,0 +1,130 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.mcp.tools; + +import com.sun.tools.attach.VirtualMachine; +import com.sun.tools.attach.VirtualMachineDescriptor; +import java.util.ArrayList; +import java.util.Collection; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Handles the list_jvms MCP tool - lists all attachable Java VMs. */ +public final class ListJvmsHandler { + private static final Logger log = LoggerFactory.getLogger(ListJvmsHandler.class); + + private ListJvmsHandler() {} + + /** Returns tool schema for MCP tools/list. */ + public static Map schema() { + Map tool = new LinkedHashMap<>(); + tool.put("name", "list_jvms"); + tool.put( + "description", + "List all attachable Java Virtual Machines on this host. " + + "Returns PID, main class, and whether BTrace is already attached (+/-). " + + "Use this to find the PID of the JVM you want to instrument."); + // No input parameters needed + Map inputSchema = new LinkedHashMap<>(); + inputSchema.put("type", "object"); + inputSchema.put("properties", new LinkedHashMap<>()); + tool.put("inputSchema", inputSchema); + return tool; + } + + /** Executes the list_jvms tool. */ + public static Map execute( + @SuppressWarnings("unused") Map arguments) { + try { + Collection> vms = listVms(); + StringBuilder sb = new StringBuilder(); + if (vms.isEmpty()) { + sb.append("No Java VMs found."); + } else { + sb.append("Discovered Java VMs:\n\n"); + for (Map vm : vms) { + sb.append( + String.format( + " PID: %s | Main Class: %s | BTrace: %s\n", + vm.get("pid"), vm.get("mainClass"), vm.get("btraceAttached"))); + } + } + return toolResult(sb.toString(), false); + } catch (Exception e) { + log.error("Failed to list JVMs", e); + return toolResult("Error listing JVMs: " + e.getMessage(), true); + } + } + + private static Collection> listVms() { + List> result = new ArrayList<>(); + try { + for (VirtualMachineDescriptor vmd : VirtualMachine.list()) { + String pid = vmd.id(); + int numericPid; + try { + numericPid = Integer.parseInt(pid); + } catch (NumberFormatException e) { + log.debug("Skipping JVM with non-numeric PID: {}", pid); + continue; + } + String mainClass = vmd.displayName(); + if (mainClass == null || mainClass.isBlank()) mainClass = ""; + Map vmInfo = new LinkedHashMap<>(); + vmInfo.put("pid", pid); + vmInfo.put("mainClass", mainClass); + vmInfo.put("btraceAttached", hasBTraceServer(numericPid) ? "attached" : "not attached"); + result.add(vmInfo); + } + } catch (Exception e) { + log.warn("Error listing VMs", e); + } + return result; + } + + private static boolean hasBTraceServer(int pid) { + com.sun.tools.attach.VirtualMachine vm = null; + try { + vm = com.sun.tools.attach.VirtualMachine.attach(String.valueOf(pid)); + return vm.getSystemProperties().containsKey("btrace.port"); + } catch (Throwable ignored) { + return false; + } finally { + if (vm != null) { + try { + vm.detach(); + } catch (Exception ignored) { + } + } + } + } + + private static Map toolResult(String text, boolean isError) { + Map content = new LinkedHashMap<>(); + content.put("type", "text"); + content.put("text", text); + List contentList = new ArrayList<>(); + contentList.add(content); + Map result = new LinkedHashMap<>(); + result.put("content", contentList); + result.put("isError", isError); + return result; + } +} diff --git a/btrace-mcp-server/src/main/java/io/btrace/mcp/tools/ListProbesHandler.java b/btrace-mcp-server/src/main/java/io/btrace/mcp/tools/ListProbesHandler.java new file mode 100644 index 000000000..89747a430 --- /dev/null +++ b/btrace-mcp-server/src/main/java/io/btrace/mcp/tools/ListProbesHandler.java @@ -0,0 +1,146 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.mcp.tools; + +import io.btrace.client.Client; +import io.btrace.core.comm.Command; +import io.btrace.core.comm.PrintableCommand; +import io.btrace.mcp.ClientManager; +import java.io.StringWriter; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Handles the list_probes MCP tool - lists active probes on a JVM. */ +public final class ListProbesHandler { + private static final Logger log = LoggerFactory.getLogger(ListProbesHandler.class); + private static final int DEFAULT_PORT = 2020; + private static final int TIMEOUT_SECONDS = 10; + + private ListProbesHandler() {} + + /** Returns tool schema for MCP tools/list. */ + public static Map schema() { + Map tool = new LinkedHashMap<>(); + tool.put("name", "list_probes"); + tool.put( + "description", + "List active BTrace probes on a running JVM. " + + "Shows which probes are currently deployed and their IDs."); + + Map properties = new LinkedHashMap<>(); + + Map pidProp = new LinkedHashMap<>(); + pidProp.put("type", "string"); + pidProp.put("description", "PID of the target JVM"); + properties.put("pid", pidProp); + + Map portProp = new LinkedHashMap<>(); + portProp.put("type", "integer"); + portProp.put("description", "BTrace agent port (default: 2020)"); + properties.put("port", portProp); + + List required = new ArrayList<>(); + required.add("pid"); + + Map inputSchema = new LinkedHashMap<>(); + inputSchema.put("type", "object"); + inputSchema.put("properties", properties); + inputSchema.put("required", required); + tool.put("inputSchema", inputSchema); + return tool; + } + + /** Executes the list_probes tool. */ + public static Map execute(Map arguments) { + String pid = getStringArg(arguments, "pid"); + int port = getIntArg(arguments, "port", DEFAULT_PORT); + + if (pid == null || pid.isEmpty()) { + return toolResult("Error: 'pid' parameter is required", true); + } + + try { + Client client = ClientManager.getClient(port); + client.attach(pid, null, "."); + + StringBuilder output = new StringBuilder(); + CountDownLatch latch = new CountDownLatch(1); + + client.connectAndListProbes( + "localhost", + cmd -> { + if (cmd instanceof PrintableCommand) { + StringWriter sw = new StringWriter(); + ((PrintableCommand) cmd).print(new java.io.PrintWriter(sw)); + output.append(sw.toString()); + } + if (cmd.getType() == Command.LIST_PROBES) { + latch.countDown(); + } + }); + + boolean done = latch.await(TIMEOUT_SECONDS, TimeUnit.SECONDS); + if (!done) { + return toolResult("Timed out waiting for probe list from PID " + pid, true); + } + + String resultText = "Active probes on PID " + pid + ":\n" + output.toString(); + return toolResult(resultText, false); + } catch (Exception e) { + log.error("Failed to list probes", e); + return toolResult("Error listing probes on PID " + pid + ": " + e.getMessage(), true); + } + } + + private static String getStringArg(Map args, String key) { + Object val = args == null ? null : args.get(key); + return val == null ? null : val.toString(); + } + + private static int getIntArg(Map args, String key, int defaultVal) { + Object val = args == null ? null : args.get(key); + if (val == null) { + return defaultVal; + } + if (val instanceof Number) { + return ((Number) val).intValue(); + } + try { + return Integer.parseInt(val.toString()); + } catch (NumberFormatException e) { + return defaultVal; + } + } + + private static Map toolResult(String text, boolean isError) { + Map content = new LinkedHashMap<>(); + content.put("type", "text"); + content.put("text", text); + List contentList = new ArrayList<>(); + contentList.add(content); + Map result = new LinkedHashMap<>(); + result.put("content", contentList); + result.put("isError", isError); + return result; + } +} diff --git a/btrace-mcp-server/src/main/java/io/btrace/mcp/tools/SendEventHandler.java b/btrace-mcp-server/src/main/java/io/btrace/mcp/tools/SendEventHandler.java new file mode 100644 index 000000000..c242a9b5d --- /dev/null +++ b/btrace-mcp-server/src/main/java/io/btrace/mcp/tools/SendEventHandler.java @@ -0,0 +1,137 @@ +/* + * Copyright (c) 2008, 2024, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.mcp.tools; + +import io.btrace.client.Client; +import io.btrace.mcp.ClientManager; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Handles the send_event MCP tool - sends an event to a running probe. */ +public final class SendEventHandler { + private static final Logger log = LoggerFactory.getLogger(SendEventHandler.class); + private static final int DEFAULT_PORT = 2020; + + private SendEventHandler() {} + + /** Returns tool schema for MCP tools/list. */ + public static Map schema() { + Map tool = new LinkedHashMap<>(); + tool.put("name", "send_event"); + tool.put( + "description", + "Send an event to a running BTrace probe. " + + "Events can trigger @OnEvent handlers in the probe script. " + + "If no event_name is specified, an unnamed event is sent."); + + Map properties = new LinkedHashMap<>(); + + Map pidProp = new LinkedHashMap<>(); + pidProp.put("type", "string"); + pidProp.put("description", "PID of the target JVM"); + properties.put("pid", pidProp); + + Map eventNameProp = new LinkedHashMap<>(); + eventNameProp.put("type", "string"); + eventNameProp.put("description", "Name of the event to send (optional)"); + properties.put("event_name", eventNameProp); + + Map portProp = new LinkedHashMap<>(); + portProp.put("type", "integer"); + portProp.put("description", "BTrace agent port (default: 2020)"); + properties.put("port", portProp); + + List required = new ArrayList<>(); + required.add("pid"); + + Map inputSchema = new LinkedHashMap<>(); + inputSchema.put("type", "object"); + inputSchema.put("properties", properties); + inputSchema.put("required", required); + tool.put("inputSchema", inputSchema); + return tool; + } + + /** Executes the send_event tool. */ + public static Map execute(Map arguments) { + String pid = getStringArg(arguments, "pid"); + String eventName = getStringArg(arguments, "event_name"); + int port = getIntArg(arguments, "port", DEFAULT_PORT); + + if (pid == null || pid.isEmpty()) { + return toolResult("Error: 'pid' parameter is required", true); + } + + try { + Client client = ClientManager.getExistingClient(pid, port); + if (client == null) { + return toolResult( + "No active BTrace session for PID " + + pid + + ". Deploy a probe first using deploy_oneliner or deploy_script.", + true); + } + + if (eventName != null && !eventName.isEmpty()) { + client.sendEvent(eventName); + return toolResult("Event '" + eventName + "' sent to PID " + pid, false); + } else { + client.sendEvent(); + return toolResult("Unnamed event sent to PID " + pid, false); + } + } catch (Exception e) { + log.error("Failed to send event", e); + return toolResult("Error sending event to PID " + pid + ": " + e.getMessage(), true); + } + } + + private static String getStringArg(Map args, String key) { + Object val = args == null ? null : args.get(key); + return val == null ? null : val.toString(); + } + + private static int getIntArg(Map args, String key, int defaultVal) { + Object val = args == null ? null : args.get(key); + if (val == null) { + return defaultVal; + } + if (val instanceof Number) { + return ((Number) val).intValue(); + } + try { + return Integer.parseInt(val.toString()); + } catch (NumberFormatException e) { + return defaultVal; + } + } + + private static Map toolResult(String text, boolean isError) { + Map content = new LinkedHashMap<>(); + content.put("type", "text"); + content.put("text", text); + List contentList = new ArrayList<>(); + contentList.add(content); + Map result = new LinkedHashMap<>(); + result.put("content", contentList); + result.put("isError", isError); + return result; + } +} diff --git a/common.gradle b/common.gradle index 462925a00..6302a93ee 100644 --- a/common.gradle +++ b/common.gradle @@ -75,6 +75,7 @@ repositories { dependencies { // Adding dependencies here will add the dependencies to each subproject. testImplementation libs.junit.jupiter + testRuntimeOnly 'org.junit.platform:junit-platform-launcher' } String mavenArtifactId = name diff --git a/docs/BTraceExtensionDevelopmentGuide.md b/docs/BTraceExtensionDevelopmentGuide.md index 6ab070fb4..d3e10c125 100644 --- a/docs/BTraceExtensionDevelopmentGuide.md +++ b/docs/BTraceExtensionDevelopmentGuide.md @@ -256,6 +256,31 @@ btraceExtension { At runtime, the agent consults this metadata to validate and enforce permissions. +## Publishing and Registry Listing + +If you want users to discover your extension through the public catalog, publish the extension artifacts to Maven Central first and then add the extension to the BTrace registry. + +The registry bootstrap currently lives in this repository at [extension-registry/](/Users/jbachorik/src/btrace/extension-registry/README.md:1) and is intended to move to a dedicated GitHub repository with GitHub Pages hosting. The canonical entry shape is documented in [ExtensionRegistry.md](ExtensionRegistry.md). + +Registry entries store a single recommended base coordinate: + +```json +{ + "id": "my-extension", + "name": "My Extension", + "description": "What it does", + "owner": "example-org", + "source_repo": "https://github.com/example-org/my-extension", + "maven": { + "groupId": "org.example", + "artifactId": "my-extension", + "version": "1.2.3" + } +} +``` + +Consumers are expected to resolve the standard BTrace extension artifacts from that base coordinate. + ## Dependency Management - Keep the API free of external library types; prefer JDK and your API classes. diff --git a/docs/BTraceTutorial.md b/docs/BTraceTutorial.md index 7099c3c25..496dd3ce5 100644 --- a/docs/BTraceTutorial.md +++ b/docs/BTraceTutorial.md @@ -1406,3 +1406,565 @@ btrace -n '/com\.example\..*/::/handle.*/ @return { time }' - Filters support simple comparisons only; no `&&`/`||` For complex scenarios — multiple probe points, state across probes, aggregations — convert to a full BTrace script. See the [Oneliner Guide](OnelinerGuide.md) for the complete syntax reference and more examples. + +--- + +## Lesson 11 — Runtime Contracts (`btrace-contracts`) + +**Extension ID:** `btrace-contracts` +**Service class:** `io.btrace.contracts.ContractService` + +`ContractService` enforces behavioral invariants at runtime — latency budgets, call-rate limits, null checks, value-range assertions — without modifying the target code. All checks are non-throwing: a violation is counted internally and the target application keeps running. It also tracks call count and average latency per user-supplied tag, so you can compare any two code paths side by side (e.g. `"cached"` vs `"direct"`, `"v1"` vs `"v2"`). + +#### API + +```java +void checkLatency(String contract, long durationNanos, long budgetNanos) +void checkCallRate(String contract, int maxPerSecond) +void assertCondition(String contract, boolean condition, String message) +void checkRange(String contract, long value, long min, long max) +void checkNotNull(String contract, Object value) +void trackCodePath(String contract, long durationNanos, String tag) +String getSummary() +boolean hasViolations() +long getTotalViolations() +``` + +#### Example — enforcing contracts on a service endpoint + +```java +import io.btrace.contracts.ContractService; +import io.btrace.core.annotations.*; +import static io.btrace.BTrace.*; + +@BTrace +public class ContractCheck { + + @Injected private static ContractService contracts; + + @OnMethod( + clazz = "com.example.RecommendationService", + method = "recommend", + location = @Location(Kind.ENTRY) + ) + public static void onEntry() { + contracts.checkCallRate("recommend/rate", 10); + } + + @OnMethod( + clazz = "com.example.RecommendationService", + method = "recommend", + location = @Location(Kind.RETURN) + ) + public static void onReturn(@Return Object result, @Duration long dur) { + contracts.checkLatency("recommend/latency", dur, 500_000_000L); + contracts.checkNotNull("recommend/non-null-result", result); + contracts.trackCodePath("recommend", dur, "impl-a"); + } + + @OnEvent + public static void report() { + println(contracts.getSummary()); + if (contracts.hasViolations()) { + println(strcat("Total violations: ", str(contracts.getTotalViolations()))); + } + } + + @OnTimer(15000) + public static void periodicReport() { report(); } +} +``` + +Copy `btrace-contracts.jar` to `$BTRACE_HOME/extensions/`, then run: + +```bash +btrace ContractCheck.java +``` + +--- + +## Lesson 12 — AI/LLM Application Observability + +Modern Java applications increasingly embed LLM inference, RAG pipelines, and on-device model execution. BTrace ships three optional extension JARs that add purpose-built services for observing these workloads without modifying application code. + +All three extensions use the standard `@Injected` mechanism: declare a field in your BTrace script, annotate it, and BTrace wires up the implementation at deploy time. No reflection, no extra threads, no allocation on the hot path. + +#### 12.1 LLM Inference Tracing (`btrace-llm-trace`) + +**Extension ID:** `btrace-llm-trace` +**Service class:** `io.btrace.llm.LlmTraceService` +**Builder class:** `io.btrace.llm.CallRecord` + +`LlmTraceService` records LLM API calls: token counts (input, output, cache-read, cache-creation), latency, streaming time-to-first-token, errors, tool calls, and embeddings. It maintains per-model statistics — call count, total tokens, and latency min/mean/max — and estimates cost using a built-in pricing table covering Claude, GPT-4o, Gemini, and other common models. The implementation is thread-safe and allocation-free on the hot path via a `ThreadLocal`-pooled builder. + +##### Simple API + +```java +void recordCall(String model, long durationNanos) +void recordCall(String model, int inputTokens, int outputTokens, long durationNanos) +CallRecord call(String model) // returns a fluent builder +void recordEmbedding(String model, int tokenCount, long durationNanos) +void recordToolUse(String model, String toolName) +void recordError(String model, String errorType, long durationNanos) +String getSummary() +double getEstimatedCostUsd() +``` + +The fluent builder lets you attach every detail of a call in one expression: + +```java +llm.call("claude-sonnet-4-20250514") + .provider("anthropic") + .inputTokens(1500) + .outputTokens(300) + .cacheReadTokens(800) + .streaming() + .timeToFirstToken(200_000_000L) + .duration(durationNanos) + .record(); +``` + +##### Full example — instrumenting LangChain4j + +The following script instruments the `ChatLanguageModel.generate()` method from the LangChain4j library and extracts token usage from the `AiMessage` response via `@Return`: + +```java +import io.btrace.llm.LlmTraceService; +import io.btrace.llm.CallRecord; +import io.btrace.core.annotations.*; +import static io.btrace.BTrace.*; + +@BTrace +public class LlmTrace { + + @Injected + private static LlmTraceService llm; + + // @Duration is available at Kind.RETURN and Kind.ERROR — no manual timestamps needed. + // @ProbeClassName identifies the concrete implementation class at the probe site. + + @OnMethod( + clazz = "+dev.langchain4j.model.chat.ChatLanguageModel", + method = "generate", + location = @Location(Kind.RETURN) + ) + public static void onChatReturn(@ProbeClassName String cls, + @Return Object response, + @Duration long durationNanos) { + // Extract token counts via BTrace field-access helpers. + // TokenUsage is nested inside Response. + Object tokenUsage = get( + field("dev.langchain4j.model.output.Response", "tokenUsage"), + response); + int inputTokens = 0; + int outputTokens = 0; + if (tokenUsage != null) { + inputTokens = (Integer) get( + field("dev.langchain4j.model.output.TokenUsage", "inputTokenCount"), + tokenUsage); + outputTokens = (Integer) get( + field("dev.langchain4j.model.output.TokenUsage", "outputTokenCount"), + tokenUsage); + } + + llm.call(cls) + .inputTokens(inputTokens) + .outputTokens(outputTokens) + .duration(durationNanos) + .record(); + } + + @OnMethod( + clazz = "+dev.langchain4j.model.chat.ChatLanguageModel", + method = "generate", + location = @Location(Kind.ERROR) + ) + public static void onChatError(@ProbeClassName String cls, + Throwable t, + @Duration long durationNanos) { + llm.recordError(cls, str(classOf(t)), durationNanos); + } + + // Print a summary on demand (btrace send event) or every 60 seconds. + @OnEvent + public static void printSummary() { + println(llm.getSummary()); + println(strcat("Estimated cost: $", + str(llm.getEstimatedCostUsd()))); + } + + @OnTimer(60000) + public static void periodicSummary() { + printSummary(); + } +} +``` + +Copy `btrace-llm-trace.jar` to `$BTRACE_HOME/extensions/` — the agent discovers and loads it automatically. Then run as normal: + +```bash +btrace LlmTrace.java +``` + +--- + +#### 12.2 RAG Pipeline Observability (`btrace-rag-quality`) + +**Extension ID:** `btrace-rag-quality` +**Service class:** `io.btrace.rag.RagQualityService` +**Builder class:** `io.btrace.rag.QueryRecord` + +`RagQualityService` tracks vector database query performance, result counts, similarity scores, and end-to-end pipeline latency broken down into retrieval and generation phases. It counts empty retrievals and chunk token sizes and works with any vector store: Pinecone, Milvus, Weaviate, Chroma, pgvector, Qdrant, and others. + +##### Simple API + +```java +void recordQuery(String source, long durationNanos) +void recordQuery(String source, int resultCount, long durationNanos) +QueryRecord query(String source) // returns a fluent builder +void recordPipeline(String pipelineName, long retrievalNanos, long generationNanos) +void recordChunk(String source, int chunkTokens) +void recordEmptyRetrieval(String source) +String getSummary() +float getAverageTopScore() +``` + +Fluent builder: + +```java +rag.query("pinecone") + .resultCount(5) + .topScore(0.92f) + .lowScore(0.71f) + .embeddingDimension(1536) + .duration(durationNanos) + .record(); +``` + +##### Full example — RAG pipeline with Pinecone and LangChain4j + +This script instruments the retrieval call to a Pinecone-backed `EmbeddingStore` and the downstream LLM call together, giving a unified view of the pipeline: + +```java +import io.btrace.rag.RagQualityService; +import io.btrace.rag.QueryRecord; +import io.btrace.llm.LlmTraceService; +import io.btrace.core.annotations.*; +import static io.btrace.BTrace.*; + +@BTrace +public class RagPipeline { + + @Injected private static RagQualityService rag; + @Injected private static LlmTraceService llm; + + // Carries the retrieval duration from the retrieval handler into the pipeline + // handler, which runs after retrieval completes on the same thread. + @TLS private static long lastRetrievalDur; + + // --- Retrieval phase --- + + @OnMethod( + clazz = "+dev.langchain4j.store.embedding.EmbeddingStore", + method = "findRelevant", + location = @Location(Kind.RETURN) + ) + public static void onRetrievalReturn(@Return Object results, + @ProbeClassName String store, + @Duration long dur) { + int count = (results != null) ? (Integer) call( + method("java.util.List", "size"), results) : 0; + + if (count == 0) { + rag.recordEmptyRetrieval(store); + } else { + rag.query(store) + .resultCount(count) + .duration(dur) + .record(); + } + lastRetrievalDur = dur; + } + + // --- Generation phase --- + + @OnMethod( + clazz = "+dev.langchain4j.chain.ConversationalRetrievalChain", + method = "execute", + location = @Location(Kind.ENTRY) + ) + public static void onPipelineEntry() { + lastRetrievalDur = 0; + } + + @OnMethod( + clazz = "+dev.langchain4j.chain.ConversationalRetrievalChain", + method = "execute", + location = @Location(Kind.RETURN) + ) + public static void onPipelineReturn(@ProbeClassName String cls, + @Duration long totalNanos) { + long generationNanos = totalNanos - lastRetrievalDur; + rag.recordPipeline(cls, lastRetrievalDur, generationNanos); + llm.recordCall(cls, generationNanos); + } + + @OnEvent + public static void printSummary() { + println("=== RAG Summary ==="); + println(rag.getSummary()); + println(strcat("Average top similarity score: ", + str(rag.getAverageTopScore()))); + println("=== LLM Summary ==="); + println(llm.getSummary()); + } + + @OnTimer(30000) + public static void periodicSummary() { + printSummary(); + } +} +``` + +Copy both extension JARs to `$BTRACE_HOME/extensions/`, then run: + +```bash +btrace RagPipeline.java +``` + +--- + +#### 12.3 GPU and Inference Observability (`btrace-gpu-bridge`) + +**Extension ID:** `btrace-gpu-bridge` +**Service class:** `io.btrace.gpu.GpuBridgeService` +**Builder class:** `io.btrace.gpu.InferenceRecord` + +`GpuBridgeService` traces the boundary between JVM code and GPU-accelerated inference runtimes: ONNX Runtime, Deep Java Library (DJL), TensorFlow Java, and Panama FFM calls into CUDA/ROCm native libraries. It tracks batch sizes, tensor dimensions, device type and ID, GPU memory allocation and deallocation, and model load times. + +##### Simple API + +```java +void recordInference(String runtime, String modelName, long durationNanos) +void recordInference(String runtime, String modelName, int batchSize, long durationNanos) +InferenceRecord inference(String runtime, String modelName) // fluent builder +void recordMemoryAlloc(String deviceType, int deviceId, long bytes) +void recordMemoryFree(String deviceType, int deviceId, long bytes) +void recordNativeCall(String library, String function, long durationNanos) +void recordModelLoad(String runtime, String modelName, long durationNanos) +String getSummary() +long getCurrentGpuMemoryBytes() +long getPeakGpuMemoryBytes() +``` + +##### Example — instrumenting ONNX Runtime + +```java +import io.btrace.gpu.GpuBridgeService; +import io.btrace.gpu.InferenceRecord; +import io.btrace.core.annotations.*; +import static io.btrace.BTrace.*; + +@BTrace +public class OnnxTrace { + + @Injected private static GpuBridgeService gpu; + + // @Self is available at Kind.RETURN, so no TLS is needed to carry the session reference. + // @Duration provides the call duration without manual timestamp diffing. + + @OnMethod( + clazz = "ai.onnxruntime.OrtSession", + method = "run", + location = @Location(Kind.RETURN) + ) + public static void onRunReturn(@Self Object session, @Duration long dur) { + String modelPath = str(get( + field("ai.onnxruntime.OrtSession", "modelPath"), session)); + gpu.inference("onnxruntime", modelPath) + .duration(dur) + .record(); + } + + @OnMethod( + clazz = "ai.onnxruntime.OrtSession", + method = "run", + location = @Location(Kind.ERROR) + ) + public static void onRunError(Throwable t, @Duration long dur) { + gpu.recordNativeCall("onnxruntime", "run/error", dur); + } + + // @Duration at Kind.RETURN gives the load duration directly; no ENTRY handler needed. + @OnMethod( + clazz = "ai.onnxruntime.OrtEnvironment", + method = "createSession", + location = @Location(Kind.RETURN) + ) + public static void onLoadReturn(@Return Object sess, @Duration long dur) { + String mp = str(get( + field("ai.onnxruntime.OrtSession", "modelPath"), sess)); + gpu.recordModelLoad("onnxruntime", mp, dur); + } + + @OnEvent + public static void printSummary() { + println(gpu.getSummary()); + println(strcat("Current GPU memory: ", + str(gpu.getCurrentGpuMemoryBytes() / (1024 * 1024)) + " MB")); + println(strcat("Peak GPU memory: ", + str(gpu.getPeakGpuMemoryBytes() / (1024 * 1024)) + " MB")); + } + + @OnTimer(60000) + public static void periodicSummary() { + printSummary(); + } +} +``` + +Copy `btrace-gpu-bridge.jar` to `$BTRACE_HOME/extensions/`, then run: + +```bash +btrace OnnxTrace.java +``` + +--- + +#### Using multiple AI extensions together + +All four extensions can be combined in a single script. Copy all the extension JARs you need to `$BTRACE_HOME/extensions/` — the agent discovers and loads them automatically: + +```bash +cp btrace-llm-trace.jar btrace-rag-quality.jar \ + btrace-gpu-bridge.jar \ + $BTRACE_HOME/extensions/ +btrace MyAIObservability.java +``` + +When using the fat agent at JVM startup, the same rule applies — extensions under `$BTRACE_HOME/extensions/` are loaded automatically: + +```bash +java -javaagent:btrace.jar=script=MyAIObservability.java -jar myapp.jar +``` + +Use `@OnEvent` for on-demand reporting triggered by `btrace send event ` and `@OnTimer` for periodic snapshots. Both can coexist in the same script. + +--- + +## Lesson 13 — BTrace MCP Server: AI Agents as Diagnosticians + +#### What is MCP? + +MCP (Model Context Protocol) is a protocol that lets AI assistants call external tools during a conversation. Instead of the AI only producing text, it can invoke structured operations — search, read a file, query a database — and incorporate the results into its response. The BTrace MCP server exposes BTrace operations as MCP tools, so an LLM client such as Claude Desktop or Claude Code can attach to running JVMs, deploy probes, read output, and clean up — all through natural language conversation. + +#### How the BTrace MCP server works + +The BTrace MCP server runs as a local subprocess on the same machine as the target JVM. The AI client starts and manages the server process; you do not need to keep a terminal open for it. When the AI calls a BTrace tool, the server forwards the request to the BTrace agent (or attaches one if none is present) and returns the result. Because the server only connects to local JVMs, and because BTrace's safety model — no loops, no allocation, no exceptions, no field assignment — still applies to every probe, the AI cannot break or hang the target application. + +#### Starting the server manually + +For testing or scripting, you can start the MCP server directly: + +```bash +java -jar btrace-mcp-server.jar +``` + +By default the server listens on `stdio` (the MCP transport used by most clients) and connects to the BTrace agent on port 2020. Override the port with: + +```bash +java -Dbtrace.port=2021 -jar btrace-mcp-server.jar +``` + +Or via the environment variable `BTRACE_PORT=2021`. + +If no BTrace agent is attached to the target JVM, the `deploy_oneliner` and `deploy_script` tools auto-attach one using the JVM's attach API — no extra step required. + +#### Client setup + +See [MCPServer.md](MCPServer.md) for per-client configuration instructions covering Claude Desktop, Claude Code, Cursor, OpenAI Codex CLI, VS Code, Windsurf, Continue.dev, and Zed. Every client follows the same pattern — point it at `btrace-mcp-server.jar` as the server command — but the configuration file location and schema differ slightly between hosts. + +#### Available tools + +| Tool | Description | +|---|---| +| `list_jvms` | List all local JVMs with their PIDs, main class, and BTrace attachment status | +| `deploy_oneliner` | Deploy a single-expression probe using oneliner syntax (auto-attaches agent) | +| `deploy_script` | Deploy a full BTrace script file to a target JVM | +| `list_probes` | List active probes on a target JVM | +| `send_event` | Send a named event to trigger `@OnEvent` handlers | +| `detach_probe` | Detach a single probe without stopping BTrace | +| `exit_probe` | Stop all BTrace probes and detach the agent from the JVM | + +#### Available prompts + +The server ships three pre-built diagnostic workflows that the AI can invoke as structured prompts: + +| Prompt | What it does | +|---|---| +| `diagnose_slow_endpoint` | Attaches timing probes to HTTP handler methods and reports latency distribution | +| `find_exception_source` | Traces exception construction and propagation across the call stack | +| `profile_method` | Measures throughput and latency of a named method over a configurable window | + +#### Example conversation + +The following exchange shows how an AI assistant uses BTrace MCP tools to diagnose a slow HTTP endpoint. Lines marked `[tool]` represent MCP tool calls made by the AI during the conversation. + +--- + +**User:** My `/recommendations` endpoint is taking over two seconds. Can you help me find out why? + +**AI:** I will start by finding which JVMs are running on this machine. + +`[tool: list_jvms]` + +``` +PID MAIN CLASS BTRACE +12345 com.example.app.ApplicationMain not attached +``` + +I found your application at PID 12345. Let me deploy a timing probe to the HTTP handler layer. + +`[tool: deploy_oneliner]` +``` +pid: 12345 +probe: /com\.example\..*/::/.*recommendations.*/i @return { time } +``` + +``` +Probe deployed. Waiting for output… +[12345] com.example.recommendation.RecommendationController.handle: 2341 ms +[12345] com.example.recommendation.RecommendationController.handle: 2289 ms +``` + +The controller itself is taking over two seconds. Let me look inside it — I will check whether the delay is in the database call or the LLM call. + +`[tool: deploy_oneliner]` +``` +pid: 12345 +probe: /com\.example\.recommendation\..*/::/.+/ @return { time } +``` + +``` +[12345] RecommendationRepository.findCandidates: 1987 ms +[12345] LlmRankingService.rank: 312 ms +``` + +The database retrieval accounts for most of the latency. The LLM ranking step is fast. Let me clean up the probes. + +`[tool: exit_probe]` +``` +pid: 12345 +``` + +``` +BTrace agent detached from PID 12345. +``` + +**AI:** The bottleneck is `RecommendationRepository.findCandidates`, which takes roughly 2 seconds out of the total 2.3 seconds. The LLM ranking step adds only about 300 ms. I recommend investigating the query or index used by `findCandidates` — adding an index on the candidate-selection columns would likely resolve most of the latency. + +--- + +#### Security note + +The BTrace MCP server only attaches to JVMs on the local machine; it cannot connect to remote processes. Every probe the AI deploys goes through BTrace's standard verifier, which enforces the same restrictions as any other BTrace script: no loops, no object allocation, no exceptions, no field writes. The AI can observe your application in detail but cannot alter its behavior or cause it to crash. diff --git a/docs/MCPServer.md b/docs/MCPServer.md new file mode 100644 index 000000000..97b6a53b8 --- /dev/null +++ b/docs/MCPServer.md @@ -0,0 +1,482 @@ +# BTrace MCP Server + +The BTrace MCP Server exposes BTrace's instrumentation capabilities through the [Model Context Protocol](https://modelcontextprotocol.io/) (MCP). This lets LLM clients — Claude Desktop, Claude Code, Cursor, and any other MCP-compatible host — instrument and diagnose running JVMs through a conversation, without leaving the AI interface. + +## How it works + +The MCP server runs as a subprocess launched by the MCP host. It speaks JSON-RPC 2.0 over stdio: the host writes requests to the process's stdin, and reads JSON responses from stdout. Log output goes to stderr and is never mixed into the protocol stream. + +When you ask the AI to "show me which methods in my service are taking more than 100 ms", it calls the appropriate BTrace MCP tools, attaches to the target JVM, deploys a probe, reads the output, and surfaces the results — all within the conversation. + +## Prerequisites + +- JDK 11 or higher to run the MCP server JAR +- The target JVM must be attachable (see [Troubleshooting](#troubleshooting)) +- `btrace-mcp-server.jar` (the fat JAR containing all dependencies) +- On JDK 8 targets: `tools.jar` from the JDK must be accessible at runtime (see [tools.jar note](#toolsjar-not-found)) + +## Quick Setup + +### Claude Desktop + +Add the following to your Claude Desktop configuration file. + +**macOS:** `~/Library/Application Support/Claude/claude_desktop_config.json` +**Windows:** `%APPDATA%\Claude\claude_desktop_config.json` + +```json +{ + "mcpServers": { + "btrace": { + "command": "java", + "args": ["-jar", "/path/to/btrace-mcp-server.jar"] + } + } +} +``` + +Replace `/path/to/btrace-mcp-server.jar` with the actual path to the JAR. After saving the file, restart Claude Desktop. You should see a tools indicator in the chat input area when BTrace tools are active. + +### Claude Code + +Run the following command in your terminal (or add it to your project's `.claude/settings.json`): + +```bash +claude mcp add btrace java -jar /path/to/btrace-mcp-server.jar +``` + +To verify the server registered correctly: + +```bash +claude mcp list +``` + +You can also add the server project-locally so that everyone working on the repository shares the same configuration: + +```json +{ + "mcpServers": { + "btrace": { + "command": "java", + "args": ["-jar", "/path/to/btrace-mcp-server.jar"] + } + } +} +``` + +Save this to `.claude/settings.json` at the root of the repository. + +### Cursor + +Cursor reads MCP server definitions from a `mcp.json` file. Use the project-local file to share the config with the team, or the user-level file to apply it globally. + +**Project-local** (committed to the repo): `.cursor/mcp.json` +**User-level** (applies to all projects): `~/.cursor/mcp.json` + +```json +{ + "mcpServers": { + "btrace": { + "command": "java", + "args": ["-jar", "/path/to/btrace-mcp-server.jar"] + } + } +} +``` + +After saving, open the Cursor command palette (`Cmd+Shift+P` / `Ctrl+Shift+P`) and run **MCP: Reload Servers**, or restart Cursor. BTrace tools appear in the agent panel once the server is active. + +### OpenAI Codex CLI + +The Codex CLI (`@openai/codex`) reads MCP server configuration from `~/.codex/config.toml` using TOML syntax: + +```toml +[mcp_servers.btrace] +command = "java" +args = ["-jar", "/path/to/btrace-mcp-server.jar"] +``` + +The key under `[mcp_servers]` becomes the server name visible in the session. Restart any active Codex session after editing the file for the change to take effect. + +If you prefer to keep the configuration per-project, create a `.codex/config.toml` at the repository root with the same content. Project-local configuration takes precedence over the user-level file. + +### VS Code (GitHub Copilot) + +VS Code exposes MCP server support through GitHub Copilot Chat's agent mode. The config file lives at `.vscode/mcp.json` within the workspace. Note that VS Code uses a slightly different schema than Claude Desktop — it requires an explicit `"type"` field: + +```json +{ + "servers": { + "btrace": { + "type": "stdio", + "command": "java", + "args": ["-jar", "/path/to/btrace-mcp-server.jar"] + } + } +} +``` + +Commit `.vscode/mcp.json` to share the setup with the rest of the team. After adding or editing the file, VS Code prompts you to reload the MCP configuration; accept the prompt or run **Developer: Reload Window**. The BTrace tools become available in Copilot Chat when you switch to agent mode (`@workspace` → agent). + +To add the server globally (not tied to a workspace), open VS Code settings (`settings.json`) and add: + +```json +"github.copilot.chat.mcp.servers": { + "btrace": { + "type": "stdio", + "command": "java", + "args": ["-jar", "/path/to/btrace-mcp-server.jar"] + } +} +``` + +### Windsurf (Codeium) + +Windsurf stores MCP configuration in `~/.codeium/windsurf/mcp_config.json`. Create the file if it does not exist: + +```json +{ + "mcpServers": { + "btrace": { + "command": "java", + "args": ["-jar", "/path/to/btrace-mcp-server.jar"] + } + } +} +``` + +Restart Windsurf after saving. The BTrace tools appear in Cascade (Windsurf's agentic panel) once the server connects. + +### Continue.dev + +Continue reads MCP servers from its configuration file. For the YAML format (Continue 0.9+), add a block under the `mcpServers` key in `~/.continue/config.yaml`: + +```yaml +mcpServers: + - name: btrace + command: java + args: + - -jar + - /path/to/btrace-mcp-server.jar +``` + +For the legacy JSON format (`~/.continue/config.json`): + +```json +{ + "mcpServers": [ + { + "name": "btrace", + "command": "java", + "args": ["-jar", "/path/to/btrace-mcp-server.jar"] + } + ] +} +``` + +Reload the Continue extension after editing (`Continue: Reload` from the command palette). BTrace tools will be available in the Continue chat panel. + +### Zed + +Zed exposes MCP tools through its context server interface. Add the server to `~/.config/zed/settings.json` under `context_servers`: + +```json +{ + "context_servers": { + "btrace": { + "command": { + "path": "java", + "args": ["-jar", "/path/to/btrace-mcp-server.jar"] + } + } + } +} +``` + +Restart Zed after saving, or use the command palette to reload context servers. BTrace tools are accessible in the Zed AI panel once the server is running. + +### Other MCP-compatible hosts + +Any host that implements the MCP specification will work. The canonical configuration pattern is: + +```json +{ + "command": "java", + "args": ["-jar", "/path/to/btrace-mcp-server.jar"], + "transport": "stdio" +} +``` + +The field names (`command`, `args`, `name`) and configuration file location vary by host. Consult your host's MCP documentation for the exact format. The BTrace MCP server itself has no host-specific requirements — it reads from stdin and writes to stdout using plain JSON-RPC 2.0. + +## Available Tools Reference + +These are the tools the MCP server exposes. Your AI client can call any of them. + +| Tool | Description | Required Parameters | Optional Parameters | +|------|-------------|---------------------|---------------------| +| `list_jvms` | Lists all attachable JVMs on the local machine | — | — | +| `deploy_oneliner` | Deploys a DTrace-style oneliner probe | `pid`, `oneliner` | `port` | +| `deploy_script` | Deploys a full BTrace Java source script | `pid`, `script` | `args`, `port` | +| `list_probes` | Lists active BTrace probes on a JVM | `pid` | `port` | +| `send_event` | Fires an event to trigger `@OnEvent` handlers | `pid` | `event_name`, `port` | +| `detach_probe` | Disconnects the client; probe keeps running | `pid` | `port` | +| `exit_probe` | Stops and removes the probe, de-instruments the JVM | `pid` | `port` | + +**Parameter notes:** + +- `pid` — integer process ID of the target JVM +- `oneliner` — a BTrace oneliner string, e.g. `"com.example.Service::process @return { print method, duration }"`; see the [Oneliner Guide](OnelinerGuide.md) for syntax +- `script` — full BTrace Java source code as a string, including the `@BTrace` annotation +- `args` — string array of arguments passed to the script's `@BTrace(unsafe=...)` or as command-line args +- `port` — BTrace agent listener port; defaults to 2020 + +## Available Prompts Reference + +The MCP server also exposes prompt templates. These generate step-by-step diagnostic instructions that the AI follows, guiding it through a structured investigation. + +| Prompt | Description | Required Arguments | Optional Arguments | +|--------|-------------|--------------------|--------------------| +| `diagnose_slow_endpoint` | Diagnose a slow HTTP endpoint; identifies where time is spent | `endpoint_class`, `endpoint_method` | `pid` | +| `find_exception_source` | Find where a specific exception originates, with stack traces | `exception_class` | `pid` | +| `profile_method` | Profile a method's latency distribution using counters | `class_name`, `method_name` | `pid` | + +## Example Session + +The following illustrates what a conversation looks like when using the BTrace MCP server. The AI calls tools automatically; you only type the natural-language requests. + +--- + +**You:** My OrderService.processOrder method seems slow. Can you find out what's happening? + +**AI:** Let me start by finding the running JVMs on this machine. + +*[calls `list_jvms`]* + +``` +PID Main Class BTrace Attached +----- --------------------------------- --------------- +18340 com.example.order.OrderApplication no +21904 org.gradle.launcher.GradleWrapperMain no +``` + +I can see your order service at PID 18340. I'll deploy a oneliner to measure the duration of `processOrder`. + +*[calls `deploy_oneliner` with pid=18340, oneliner="com.example.order.OrderService::processOrder @return { print method, duration }"]* + +``` +Probe deployed. Output: +processOrder: 342ms +processOrder: 289ms +processOrder: 1204ms +processOrder: 301ms +``` + +The method is occasionally taking over a second. Let me drill into all methods inside `OrderService` to find the bottleneck. + +*[calls `deploy_oneliner` with pid=18340, oneliner="com.example.order.OrderService::/.*/ @return if duration>50ms { print method, duration }"]* + +``` +Probe deployed. Output: +fetchInventory: 1150ms +validatePayment: 12ms +``` + +`fetchInventory` is the culprit — it accounts for almost all of the slow calls. That looks like a database or external service call worth investigating. I'll remove the probe now. + +*[calls `exit_probe` with pid=18340]* + +``` +Probe removed. +``` + +--- + +The AI handled every tool call. You described the goal; BTrace provided the data. + +## Typical Workflow + +For most diagnostic tasks, the sequence is: + +1. **`list_jvms`** — identify which process you want to inspect and get its PID. +2. **`deploy_oneliner`** — deploy a quick measurement for the method or class you suspect. +3. Read the output that streams back through the conversation. +4. **`deploy_script`** — if you need more detail, deploy a full BTrace script with state, aggregation, or `@OnEvent` triggers. +5. **`send_event`** (optional) — trigger `@OnEvent` handlers in your script to dump summaries or flush aggregated state. +6. **`list_probes`** (optional) — verify what probes are active if you lose track. +7. **`exit_probe`** — remove the probe and restore the JVM to its original state. + +Use `detach_probe` instead of `exit_probe` if you want the probe to keep collecting data after the client disconnects. You can reattach to a running probe later. + +## Writing Scripts for the MCP Server + +Scripts passed to `deploy_script` are standard BTrace Java source. The `@BTrace` annotation and imports are required. + +### Minimal example + +```java +import io.btrace.core.annotations.*; +import static io.btrace.core.BTraceUtils.*; + +@BTrace +public class MethodTimer { + @OnMethod( + clazz = "com.example.Service", + method = "process", + location = @Location(Kind.RETURN) + ) + public static void onReturn(@Duration long duration) { + println("process: " + (duration / 1_000_000) + "ms"); + } +} +``` + +### Script with event-triggered summary + +```java +import io.btrace.core.annotations.*; +import static io.btrace.core.BTraceUtils.*; + +@BTrace +public class CallCounter { + private static long calls; + private static long errors; + + @OnMethod(clazz = "com.example.Service", method = "process") + public static void onEntry() { + calls++; + } + + @OnMethod( + clazz = "com.example.Service", + method = "process", + location = @Location(Kind.ERROR) + ) + public static void onError() { + errors++; + } + + @OnEvent + public static void printSummary() { + println("calls=" + calls + " errors=" + errors); + } +} +``` + +After deploying this script, call `send_event` to trigger `printSummary` at any time. + +See the [BTrace Tutorial](BTraceTutorial.md) and [Quick Reference](QuickReference.md) for the full annotation set and built-in functions. + +## Security Considerations + +**Local attachment only.** The BTrace MCP server uses the JVM Attach API (`VirtualMachine.list()` and `VirtualMachine.attach()`), which only works for JVMs running on the same machine under the same OS user. It cannot attach to remote JVMs or JVMs owned by a different user. There is no network exposure: the MCP transport is stdio, so there is no listening socket to exploit. + +**BTrace safety model applies.** All scripts deployed through the MCP server go through BTrace's standard compile-time verifier, which rejects: + +- Loops and recursion +- Object allocation with `new` +- `throw` statements +- Field assignments + +These restrictions prevent BTrace scripts from destabilizing the target JVM. A probe that violates these rules will fail to compile and will not be deployed. + +**No persistent access.** `exit_probe` removes all instrumentation and returns the JVM to its original state. The agent is unloaded when all probes are removed. + +**JDK 21+ dynamic agent loading warning.** From JDK 21 onward (JEP 451), dynamically attaching an agent prints a warning to the target JVM's stderr. The probe still works. To suppress the warning, start the target JVM with: + +``` +-XX:+EnableDynamicAgentLoading +``` + +## Troubleshooting + +### tools.jar not found + +**Symptom:** The server starts but `list_jvms` returns an empty list or fails with a class-not-found error related to `com.sun.tools.attach`. + +**Cause:** The JVM running the MCP server cannot find the Attach API classes. On JDK 8 these live in `tools.jar`; from JDK 9 onward the Attach API is in the `jdk.attach` module and is always available. + +**Solution:** Make sure you are running the MCP server with a full JDK (not a JRE), JDK 11 or higher: + +```bash +java -version # should report 11 or above, from a JDK distribution +java -jar btrace-mcp-server.jar +``` + +If you must use JDK 8, add `tools.jar` to the classpath explicitly: + +```bash +java -cp "$JAVA_HOME/lib/tools.jar:btrace-mcp-server.jar" io.btrace.mcp.BTraceMcpServer +``` + +### Port already in use + +**Symptom:** `deploy_oneliner` or `deploy_script` fails with a connection error mentioning port 2020. + +**Cause:** Another BTrace client or agent is already listening or connected on the default port. + +**Solution:** Pass a different `port` value to the deploy tools, and make sure the BTrace agent on the target JVM is started with a matching port. For example, if the agent was loaded with `-javaagent:btrace.jar=port=2021`, use `port=2021` in every MCP tool call for that JVM. + +### JVM not attachable + +**Symptom:** `list_jvms` shows the process but `deploy_oneliner` fails with an attach exception. + +**Common causes and fixes:** + +| Cause | Fix | +|-------|-----| +| User mismatch: MCP server runs as a different OS user than the target JVM | Start the MCP server (and therefore the MCP host) as the same user as the target JVM | +| Target JVM started with `-XX:+DisableAttachMechanism` | Remove that flag and restart the application | +| JDK 21+: dynamic agent loading disabled | Add `-XX:+EnableDynamicAgentLoading` to the target JVM's startup flags | +| Target is a JRE, not a JDK | Run the target application with a full JDK | + +See [JVM Attachment Issues](Troubleshooting.md#jvm-attachment-issues) in the main Troubleshooting Guide for a more complete list. + +### MCP server not appearing in the client + +**Symptom:** After adding the server configuration, tools do not appear in the AI client. + +**Steps:** + +1. Verify the JAR path is correct and the file exists. +2. Test the server manually: `java -jar btrace-mcp-server.jar` should start without error and wait on stdin. +3. Check that you are using JDK 11+. +4. Restart the MCP host (Claude Desktop, Cursor, etc.) after changing the config file. +5. Check the MCP host's logs for subprocess errors (Claude Desktop logs to `~/Library/Logs/Claude/` on macOS). + +### Probe output not visible in the conversation + +**Symptom:** The probe deploys successfully but no output appears. + +**Cause:** The target method may not be called during the observation window, or the class/method pattern does not match. + +**Steps:** + +1. Use `list_probes` to confirm the probe is active. +2. Trigger the target code path in the application (send a request, run the workflow, etc.). +3. Verify the class and method names are correct — BTrace matching is case-sensitive and requires fully qualified class names. +4. Narrow the pattern: if you used a wildcard, try an exact class name first. +5. Check whether the class is loaded by a non-system classloader; see [Class/Method Not Found](Troubleshooting.md#classmethod-not-found) in the Troubleshooting Guide. + +### Script verification failure + +**Symptom:** `deploy_script` returns an error mentioning "unsafe" or "verification failed". + +**Cause:** The script violates one of BTrace's safety restrictions (loop, allocation, throw, field write). + +**Fix:** Review the script against the [BTrace Script Restrictions](GettingStarted.md#btrace-script-restrictions) and remove the offending construct. For cases where safety restrictions are intentionally relaxed, pass `unsafe=true` in the `@BTrace` annotation and include `"unsafe"` in the `args` array: + +```java +@BTrace(unsafe = true) +public class UnsafeScript { ... } +``` + +Note that unsafe scripts require the BTrace agent to be started in unsafe mode as well. + +## Related Documentation + +- [Getting Started Guide](GettingStarted.md) — installation, first script, deployment modes +- [Oneliner Guide](OnelinerGuide.md) — oneliner syntax reference for `deploy_oneliner` +- [Quick Reference](QuickReference.md) — all annotations, locations, and built-in functions +- [BTrace Tutorial](BTraceTutorial.md) — comprehensive guide to writing BTrace scripts +- [Troubleshooting Guide](Troubleshooting.md) — common errors and their solutions diff --git a/docs/README.md b/docs/README.md index 5c9270ffa..f1fe1ade5 100644 --- a/docs/README.md +++ b/docs/README.md @@ -18,6 +18,8 @@ Get up and running in 5 minutes with installation, your first script, and common | **[BTrace Tutorial](BTraceTutorial.md)** | Comprehensive lessons covering all features | All users, in-depth learning | | **[Troubleshooting Guide](Troubleshooting.md)** | Common errors, debugging, performance, compatibility | Problem-solving, debugging | | **[FAQ](FAQ.md)** | Common questions, best practices, comparisons | All users, decision-making | +| **[MCP Server](MCPServer.md)** | Using BTrace from AI clients via the Model Context Protocol | AI-assisted debugging, LLM integrations | +| **[Extension Registry](ExtensionRegistry.md)** | JSON catalog and GitHub workflow for discovering published extensions | Extension authors, platform maintainers | ## Learning Paths @@ -52,6 +54,12 @@ Tip: Want latency histograms fast? See [Quick Start: Histogram Metrics Extension 6. **Level Filtering** → [Quick Reference: @Level](QuickReference.md#level) 7. **Extensions Architecture** → [Extension invokedynamic Bridge](architecture/ExtensionInvokeDynamicBridge.md) +### I'm Working on AI/LLM Applications +1. **Observe LLM API calls** → [Tutorial Lesson 12.1](BTraceTutorial.md#121-llm-inference-tracing-btrace-llm-trace) — token counts, latency, cost tracking +2. **Observe RAG pipelines** → [Tutorial Lesson 12.2](BTraceTutorial.md#122-rag-pipeline-observability-btrace-rag-quality) — vector DB queries, similarity scores +3. **Observe on-device inference** → [Tutorial Lesson 12.3](BTraceTutorial.md#123-gpu-and-inference-observability-btrace-gpu-bridge) — ONNX, DJL, TensorFlow +4. **Use BTrace from an AI agent** → [MCP Server](MCPServer.md) for Claude Desktop, Claude Code, Cursor + ## Documentation by Topic ### Core Features @@ -66,6 +74,13 @@ Tip: Want latency histograms fast? See [Quick Start: Histogram Metrics Extension - **Level Control** → [Quick Reference: @Level](QuickReference.md#level) - **Aggregations** → [Quick Reference: Aggregation Functions](QuickReference.md#aggregation-functions) - **Periodic Events** → [Quick Reference: @OnTimer](QuickReference.md#ontimer), [@PeriodicEvent](QuickReference.md#periodicevent) +- **Runtime Contracts** → [Tutorial Lesson 11](BTraceTutorial.md#lesson-11--runtime-contracts-btrace-contracts) — latency budgets, call-rate limits, assertions, tagged path profiling + +### AI/LLM Observability +- **LLM Inference Tracing** → [Tutorial Lesson 12.1](BTraceTutorial.md#121-llm-inference-tracing-btrace-llm-trace) — token counts, latency, cost, streaming TTFT +- **RAG Pipeline Observability** → [Tutorial Lesson 12.2](BTraceTutorial.md#122-rag-pipeline-observability-btrace-rag-quality) — vector DB queries, similarity scores, empty retrievals +- **GPU / Inference Observability** → [Tutorial Lesson 12.3](BTraceTutorial.md#123-gpu-and-inference-observability-btrace-gpu-bridge) — ONNX Runtime, DJL, TensorFlow, Panama FFM +- **AI Agents (MCP)** → [MCP Server](MCPServer.md) — connect Claude Desktop, Claude Code, or Cursor to live JVMs ### Deployment & Operations - **Installation** → [Getting Started: Installation](GettingStarted.md#installation) @@ -88,11 +103,14 @@ Tip: Want latency histograms fast? See [Quick Start: Histogram Metrics Extension - **Third-Party Libraries** → [FAQ: Third-Party](FAQ.md#how-do-i-trace-methods-from-third-party-libraries) - **JMX Export** → [Quick Reference: @Export](QuickReference.md#export), [FAQ: Monitoring Integration](FAQ.md#can-i-integrate-btrace-with-monitoring-systems) - **Service Mesh** → [FAQ: Service Mesh](FAQ.md#does-btrace-work-with-service-meshes-istiolinkerd) +- **AI / LLM Clients (MCP)** → [MCP Server](MCPServer.md) for Claude Desktop, Claude Code, Cursor +- **Published Extensions** → [Extension Registry](ExtensionRegistry.md) for the catalog format and PR flow ### Architecture - **Masked JAR** → [Masked JAR Architecture](architecture/MaskedJarArchitecture.md) — single-JAR distribution with classdata masking - **v2 Binary Protocol** → [Version 2 Protocol Architecture](architecture/Version2ProtocolArchitecture.md) — custom binary serialization - **Extension Framework** → [Extension invokedynamic Bridge](architecture/ExtensionInvokeDynamicBridge.md), [Extension Configuration](architecture/ExtensionConfiguration.md), [Extension Manifest](architecture/ExtensionManifestFormat.md), [Extension Storage](architecture/ExtensionStorageDesign.md) +- **Extension Catalog** → [Extension Registry](ExtensionRegistry.md) for the public registry contract - **Instrumentation** → [BTrace Instrumentation Analysis](architecture/BTraceInstrAnalysis.md) ## Sample Scripts diff --git a/gradle/wrapper/gradle-wrapper.properties b/gradle/wrapper/gradle-wrapper.properties index aaaabb3cb..5dd3c0121 100644 --- a/gradle/wrapper/gradle-wrapper.properties +++ b/gradle/wrapper/gradle-wrapper.properties @@ -1,6 +1,6 @@ distributionBase=GRADLE_USER_HOME distributionPath=wrapper/dists -distributionUrl=https\://services.gradle.org/distributions/gradle-8.14.4-bin.zip +distributionUrl=https\://services.gradle.org/distributions/gradle-9.5.1-bin.zip networkTimeout=10000 validateDistributionUrl=true zipStoreBase=GRADLE_USER_HOME diff --git a/internal/superpowers/plans/2026-05-17-extension-tck-plan-1-structural.md b/internal/superpowers/plans/2026-05-17-extension-tck-plan-1-structural.md new file mode 100644 index 000000000..9333997aa --- /dev/null +++ b/internal/superpowers/plans/2026-05-17-extension-tck-plan-1-structural.md @@ -0,0 +1,2328 @@ +# Extension TCK — Part 1: Structural Suite Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Build `btrace-ext-validator` (shared validation logic) and the structural suite of `btrace-tck` (manifest, layout, service API, and API/impl partition checks), producing a working standalone JAR that can certify an extension ZIP without launching a JVM. + +**Architecture:** Two new Gradle modules in the mono-repo: `btrace-ext-validator` (pure Java, ASM-based bytecode analysis) provides all validation logic; `btrace-tck` depends on it, wraps checks in a suite/engine abstraction, and writes JUnit XML + HTML + stdout reports. The Gradle plugin is updated to delegate `validateServiceApis` to `btrace-ext-validator` so the rule logic has a single home. + +**Tech Stack:** Java 11 (toolchain), ASM 9.9.1, JUnit 5.11.4, Jackson 2.18.2, picocli 4.7.6 + +--- + +## File Map + +### New: `btrace-ext-validator/` +| File | Responsibility | +|------|---------------| +| `build.gradle` | Module build config | +| `src/main/java/io/btrace/ext/validator/ValidationSeverity.java` | ERROR / WARN / INFO enum | +| `src/main/java/io/btrace/ext/validator/ValidationResult.java` | Single rule finding | +| `src/main/java/io/btrace/ext/validator/ManifestValidator.java` | BTRACE-MF-* manifest attribute checks | +| `src/main/java/io/btrace/ext/validator/ArtifactLayoutValidator.java` | BTRACE-LO-* ZIP structure checks | +| `src/main/java/io/btrace/ext/validator/ServiceApiValidator.java` | BTRACE-EXT-001..041 bytecode checks | +| `src/main/java/io/btrace/ext/validator/ApiImplPartitionValidator.java` | BTRACE-PT-* impl-leak-into-API checks | +| `src/test/java/io/btrace/ext/validator/*Test.java` | Per-validator unit tests | + +### New: `btrace-tck/` +| File | Responsibility | +|------|---------------| +| `build.gradle` | Module build config | +| `src/main/java/io/btrace/tck/TckStatus.java` | PASS / FAIL / SKIP | +| `src/main/java/io/btrace/tck/TckResult.java` | One check outcome | +| `src/main/java/io/btrace/tck/TckSuiteResult.java` | All outcomes for one suite | +| `src/main/java/io/btrace/tck/TckInput.java` | Extension ZIP + btrace home + config path | +| `src/main/java/io/btrace/tck/TckConfig.java` | Parsed tck-config.yaml | +| `src/main/java/io/btrace/tck/TckEngine.java` | Orchestrates suites, returns results | +| `src/main/java/io/btrace/tck/suite/StructuralSuite.java` | Runs all structural checks | +| `src/main/java/io/btrace/tck/check/structural/ManifestCheck.java` | Delegates to ManifestValidator | +| `src/main/java/io/btrace/tck/check/structural/ArtifactLayoutCheck.java` | Delegates to ArtifactLayoutValidator | +| `src/main/java/io/btrace/tck/check/structural/ServiceApiCheck.java` | Delegates to ServiceApiValidator | +| `src/main/java/io/btrace/tck/check/structural/ApiImplPartitionCheck.java` | Delegates to ApiImplPartitionValidator | +| `src/main/java/io/btrace/tck/report/StdoutReporter.java` | Real-time [PASS]/[FAIL] lines | +| `src/main/java/io/btrace/tck/report/JUnitXmlReporter.java` | tck-results.xml | +| `src/main/java/io/btrace/tck/report/HtmlReporter.java` | tck-report.html | +| `src/main/java/io/btrace/tck/cli/TckMain.java` | Standalone JAR entry point | +| `src/test/resources/fixtures/good-extension/` | Valid extension ZIP fixture | +| `src/test/resources/fixtures/bad-manifest/` | Missing required manifest attributes | +| `src/test/resources/fixtures/bad-api/` | BTRACE-EXT-001 violation | +| `src/test/resources/fixtures/bad-partition/` | Impl class in api.jar | + +### Modified +| File | Change | +|------|--------| +| `settings.gradle` | Add `include 'btrace-ext-validator'`, `include 'btrace-tck'` | +| `btrace-gradle-plugin/src/main/groovy/io/btrace/gradle/BTraceExtensionPlugin.groovy` | Delegate `validateServiceApis` to `ServiceApiValidator` | + +--- + +## Task 1: btrace-ext-validator module skeleton + +**Files:** +- Create: `btrace-ext-validator/build.gradle` +- Create: `btrace-ext-validator/src/main/java/io/btrace/ext/validator/ValidationSeverity.java` +- Create: `btrace-ext-validator/src/main/java/io/btrace/ext/validator/ValidationResult.java` +- Modify: `settings.gradle` + +- [ ] **Step 1: Add module to settings.gradle** + +Open `settings.gradle` and add after the last `include` line: +```groovy +include 'btrace-ext-validator' +``` + +- [ ] **Step 2: Create build.gradle** + +Create `btrace-ext-validator/build.gradle`: +```groovy +apply from: rootProject.file('common.gradle') + +java { + toolchain { languageVersion = JavaLanguageVersion.of(11) } +} + +dependencies { + implementation 'org.ow2.asm:asm:9.9.1' + implementation 'org.ow2.asm:asm-tree:9.9.1' + testImplementation platform('org.junit:junit-bom:5.11.4') + testImplementation 'org.junit.jupiter:junit-jupiter' + testRuntimeOnly 'org.junit.platform:junit-platform-launcher' +} + +test { useJUnitPlatform() } +``` + +- [ ] **Step 3: Write failing test** + +Create `btrace-ext-validator/src/test/java/io/btrace/ext/validator/ValidationResultTest.java`: +```java +package io.btrace.ext.validator; + +import org.junit.jupiter.api.Test; +import static org.junit.jupiter.api.Assertions.*; + +class ValidationResultTest { + @Test + void errorResultIsError() { + var r = new ValidationResult("BTRACE-MF-001", ValidationSeverity.ERROR, "Id missing", "MANIFEST.MF"); + assertTrue(r.isError()); + assertEquals("BTRACE-MF-001", r.getRuleCode()); + assertEquals("MANIFEST.MF", r.getArtifact()); + } + + @Test + void warnResultIsNotError() { + var r = new ValidationResult("BTRACE-MF-004", ValidationSeverity.WARN, "Name blank", "MANIFEST.MF"); + assertFalse(r.isError()); + } + + @Test + void nullArtifactBecomesEmpty() { + var r = new ValidationResult("X", ValidationSeverity.INFO, "msg", null); + assertEquals("", r.getArtifact()); + } +} +``` + +- [ ] **Step 4: Run test — expect compilation failure** + +```bash +./gradlew :btrace-ext-validator:test 2>&1 | tail -20 +``` +Expected: compilation error (classes not yet created). + +- [ ] **Step 5: Create ValidationSeverity.java** + +```java +package io.btrace.ext.validator; + +public enum ValidationSeverity { ERROR, WARN, INFO } +``` + +- [ ] **Step 6: Create ValidationResult.java** + +```java +package io.btrace.ext.validator; + +import java.util.Objects; + +public final class ValidationResult { + private final String ruleCode; + private final ValidationSeverity severity; + private final String message; + private final String artifact; + + public ValidationResult(String ruleCode, ValidationSeverity severity, String message, String artifact) { + this.ruleCode = Objects.requireNonNull(ruleCode, "ruleCode"); + this.severity = Objects.requireNonNull(severity, "severity"); + this.message = Objects.requireNonNull(message, "message"); + this.artifact = artifact != null ? artifact : ""; + } + + public String getRuleCode() { return ruleCode; } + public ValidationSeverity getSeverity() { return severity; } + public String getMessage() { return message; } + public String getArtifact() { return artifact; } + public boolean isError() { return severity == ValidationSeverity.ERROR; } + + @Override public String toString() { + return "[" + severity + "] " + ruleCode + " " + message + + (artifact.isEmpty() ? "" : " (" + artifact + ")"); + } +} +``` + +- [ ] **Step 7: Run tests — expect PASS** + +```bash +./gradlew :btrace-ext-validator:test +``` +Expected: `BUILD SUCCESSFUL`, 3 tests passed. + +- [ ] **Step 8: Commit** + +```bash +git add btrace-ext-validator/ settings.gradle +git commit -m "feat(tck): add btrace-ext-validator module skeleton" +``` + +--- + +## Task 2: ManifestValidator + +**Files:** +- Create: `btrace-ext-validator/src/main/java/io/btrace/ext/validator/ManifestValidator.java` +- Create: `btrace-ext-validator/src/test/java/io/btrace/ext/validator/ManifestValidatorTest.java` + +Manifest rule codes: +- `BTRACE-MF-001` ERROR — `BTrace-Extension-Id` missing or blank +- `BTRACE-MF-002` ERROR — `BTrace-Extension-Id` invalid format (must match `[a-zA-Z0-9][a-zA-Z0-9._-]*`) +- `BTRACE-MF-003` ERROR — `BTrace-Extension-Version` missing or blank +- `BTRACE-MF-004` WARN — `BTrace-Extension-Name` missing or blank +- `BTRACE-MF-005` WARN — `BTrace-Extension-Description` missing or blank +- `BTRACE-MF-006` ERROR — `BTrace-API-Version` missing or blank +- `BTRACE-MF-007` ERROR — `BTrace-API-Version` invalid (must match `\d+\.\d+(\.\d+)?[+]?` or range like `[2.3,3.0)`) +- `BTRACE-MF-008` WARN — `BTrace-Java-Version` missing (defaults to `8+`, warn only) +- `BTRACE-MF-009` ERROR — `BTrace-Extension-Services` missing or blank +- `BTRACE-MF-010` ERROR — `BTrace-Extension-Impl` missing or blank + +- [ ] **Step 1: Write failing tests** + +Create `btrace-ext-validator/src/test/java/io/btrace/ext/validator/ManifestValidatorTest.java`: +```java +package io.btrace.ext.validator; + +import org.junit.jupiter.api.Test; +import java.util.jar.Attributes; +import java.util.jar.Manifest; +import java.util.List; +import static org.junit.jupiter.api.Assertions.*; + +class ManifestValidatorTest { + private static Manifest validManifest() { + var m = new Manifest(); + var a = m.getMainAttributes(); + a.put(Attributes.Name.MANIFEST_VERSION, "1.0"); + a.putValue("BTrace-Extension-Id", "io.example.my-ext"); + a.putValue("BTrace-Extension-Version", "1.0.0"); + a.putValue("BTrace-Extension-Name", "My Extension"); + a.putValue("BTrace-Extension-Description", "Does things"); + a.putValue("BTrace-API-Version", "3.0+"); + a.putValue("BTrace-Java-Version", "8+"); + a.putValue("BTrace-Extension-Services","io.example.MyService"); + a.putValue("BTrace-Extension-Impl", "my-ext-1.0.0-impl.jar"); + return m; + } + + @Test void validManifestProducesNoErrors() { + var results = new ManifestValidator().validate(validManifest()); + assertTrue(results.stream().noneMatch(ValidationResult::isError), + () -> "Unexpected errors: " + results); + } + + @Test void missingIdIsError() { + var m = validManifest(); + m.getMainAttributes().remove(new Attributes.Name("BTrace-Extension-Id")); + var results = new ManifestValidator().validate(m); + assertTrue(results.stream().anyMatch(r -> r.getRuleCode().equals("BTRACE-MF-001") && r.isError())); + } + + @Test void invalidIdFormatIsError() { + var m = validManifest(); + m.getMainAttributes().putValue("BTrace-Extension-Id", "../etc/passwd"); + var results = new ManifestValidator().validate(m); + assertTrue(results.stream().anyMatch(r -> r.getRuleCode().equals("BTRACE-MF-002") && r.isError())); + } + + @Test void missingVersionIsError() { + var m = validManifest(); + m.getMainAttributes().remove(new Attributes.Name("BTrace-Extension-Version")); + var results = new ManifestValidator().validate(m); + assertTrue(results.stream().anyMatch(r -> r.getRuleCode().equals("BTRACE-MF-003") && r.isError())); + } + + @Test void missingServicesIsError() { + var m = validManifest(); + m.getMainAttributes().remove(new Attributes.Name("BTrace-Extension-Services")); + var results = new ManifestValidator().validate(m); + assertTrue(results.stream().anyMatch(r -> r.getRuleCode().equals("BTRACE-MF-009") && r.isError())); + } + + @Test void missingImplIsError() { + var m = validManifest(); + m.getMainAttributes().remove(new Attributes.Name("BTrace-Extension-Impl")); + var results = new ManifestValidator().validate(m); + assertTrue(results.stream().anyMatch(r -> r.getRuleCode().equals("BTRACE-MF-010") && r.isError())); + } + + @Test void missingNameIsWarnOnly() { + var m = validManifest(); + m.getMainAttributes().remove(new Attributes.Name("BTrace-Extension-Name")); + var results = new ManifestValidator().validate(m); + assertTrue(results.stream().anyMatch(r -> r.getRuleCode().equals("BTRACE-MF-004") && !r.isError())); + assertTrue(results.stream().noneMatch(ValidationResult::isError)); + } +} +``` + +- [ ] **Step 2: Run tests — expect compilation failure** + +```bash +./gradlew :btrace-ext-validator:test 2>&1 | tail -10 +``` +Expected: `ManifestValidator` not found. + +- [ ] **Step 3: Implement ManifestValidator** + +Create `btrace-ext-validator/src/main/java/io/btrace/ext/validator/ManifestValidator.java`: +```java +package io.btrace.ext.validator; + +import java.util.ArrayList; +import java.util.List; +import java.util.jar.Attributes; +import java.util.jar.Manifest; +import java.util.regex.Pattern; + +public final class ManifestValidator { + private static final Pattern VALID_ID = + Pattern.compile("[a-zA-Z0-9][a-zA-Z0-9._-]*"); + private static final Pattern VALID_API_VERSION = + Pattern.compile("(\\d+\\.\\d+(\\.\\d+)?[+]?|[\\[\\(].*[\\]\\)])"); + + public List validate(Manifest manifest) { + var results = new ArrayList(); + var attrs = manifest.getMainAttributes(); + + check(results, attrs, "BTrace-Extension-Id", "BTRACE-MF-001", "BTRACE-MF-002", + ValidationSeverity.ERROR, VALID_ID, true); + requirePresent(results, attrs, "BTrace-Extension-Version", "BTRACE-MF-003", ValidationSeverity.ERROR); + requirePresent(results, attrs, "BTrace-Extension-Name", "BTRACE-MF-004", ValidationSeverity.WARN); + requirePresent(results, attrs, "BTrace-Extension-Description", "BTRACE-MF-005", ValidationSeverity.WARN); + check(results, attrs, "BTrace-API-Version", "BTRACE-MF-006", "BTRACE-MF-007", + ValidationSeverity.ERROR, VALID_API_VERSION, true); + requirePresent(results, attrs, "BTrace-Java-Version", "BTRACE-MF-008", ValidationSeverity.WARN); + requirePresent(results, attrs, "BTrace-Extension-Services", "BTRACE-MF-009", ValidationSeverity.ERROR); + requirePresent(results, attrs, "BTrace-Extension-Impl", "BTRACE-MF-010", ValidationSeverity.ERROR); + + return results; + } + + private void requirePresent(List out, Attributes attrs, + String key, String code, ValidationSeverity sev) { + String v = attrs.getValue(key); + if (v == null || v.isBlank()) { + out.add(new ValidationResult(code, sev, + key + " is missing or blank", "MANIFEST.MF")); + } + } + + private void check(List out, Attributes attrs, String key, + String missingCode, String formatCode, ValidationSeverity sev, + Pattern pattern, boolean required) { + String v = attrs.getValue(key); + if (v == null || v.isBlank()) { + if (required) out.add(new ValidationResult(missingCode, sev, + key + " is missing or blank", "MANIFEST.MF")); + return; + } + if (!pattern.matcher(v).matches()) { + out.add(new ValidationResult(formatCode, sev, + key + " has invalid format: " + v, "MANIFEST.MF")); + } + } +} +``` + +- [ ] **Step 4: Run tests — expect PASS** + +```bash +./gradlew :btrace-ext-validator:test +``` +Expected: `BUILD SUCCESSFUL`, 8 tests passed. + +- [ ] **Step 5: Commit** + +```bash +git add btrace-ext-validator/ +git commit -m "feat(tck): add ManifestValidator with BTRACE-MF-* rules" +``` + +--- + +## Task 3: ArtifactLayoutValidator + +**Files:** +- Create: `btrace-ext-validator/src/main/java/io/btrace/ext/validator/ArtifactLayoutValidator.java` +- Create: `btrace-ext-validator/src/test/java/io/btrace/ext/validator/ArtifactLayoutValidatorTest.java` + +Rule codes: +- `BTRACE-LO-001` ERROR — extension path does not exist or cannot be opened as ZIP +- `BTRACE-LO-002` ERROR — no file matching `*-api.jar` found in ZIP +- `BTRACE-LO-003` ERROR — api.jar has no readable MANIFEST.MF +- `BTRACE-LO-004` ERROR — `BTrace-Extension-Impl` in api.jar manifest doesn't match any entry in ZIP +- `BTRACE-LO-005` WARN — impl JAR filename doesn't follow `{name}-{version}-impl.jar` convention + +- [ ] **Step 1: Write failing test** + +Create `btrace-ext-validator/src/test/java/io/btrace/ext/validator/ArtifactLayoutValidatorTest.java`: +```java +package io.btrace.ext.validator; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import java.io.*; +import java.nio.file.*; +import java.util.jar.*; +import java.util.zip.*; +import static org.junit.jupiter.api.Assertions.*; + +class ArtifactLayoutValidatorTest { + @TempDir Path tmp; + + private Path buildValidZip() throws Exception { + // Build api.jar + var apiJar = tmp.resolve("my-ext-1.0.0-api.jar"); + try (var jos = new JarOutputStream(new FileOutputStream(apiJar.toFile()))) { + var mf = new Manifest(); + mf.getMainAttributes().put(Attributes.Name.MANIFEST_VERSION, "1.0"); + mf.getMainAttributes().putValue("BTrace-Extension-Id", "io.example.my-ext"); + mf.getMainAttributes().putValue("BTrace-Extension-Impl", "my-ext-1.0.0-impl.jar"); + var mfEntry = new JarEntry("META-INF/MANIFEST.MF"); + jos.putNextEntry(mfEntry); + mf.write(jos); + jos.closeEntry(); + } + var implJar = tmp.resolve("my-ext-1.0.0-impl.jar"); + implJar.toFile().createNewFile(); + + var zip = tmp.resolve("my-ext-1.0.0-extension.zip"); + try (var zos = new ZipOutputStream(new FileOutputStream(zip.toFile()))) { + for (Path p : List.of(apiJar, implJar)) { + zos.putNextEntry(new ZipEntry(p.getFileName().toString())); + Files.copy(p, zos); + zos.closeEntry(); + } + } + return zip; + } + + @Test void validZipProducesNoErrors() throws Exception { + var results = new ArtifactLayoutValidator().validate(buildValidZip()); + assertTrue(results.stream().noneMatch(ValidationResult::isError), + () -> "Unexpected errors: " + results); + } + + @Test void missingApiJarIsError() throws Exception { + var zip = tmp.resolve("empty.zip"); + new ZipOutputStream(new FileOutputStream(zip.toFile())).close(); + var results = new ArtifactLayoutValidator().validate(zip); + assertTrue(results.stream().anyMatch(r -> r.getRuleCode().equals("BTRACE-LO-002") && r.isError())); + } + + @Test void missingImplJarIsError() throws Exception { + var apiJar = tmp.resolve("my-ext-1.0.0-api.jar"); + try (var jos = new JarOutputStream(new FileOutputStream(apiJar.toFile()))) { + var mf = new Manifest(); + mf.getMainAttributes().put(Attributes.Name.MANIFEST_VERSION, "1.0"); + mf.getMainAttributes().putValue("BTrace-Extension-Impl", "my-ext-1.0.0-impl.jar"); + jos.putNextEntry(new JarEntry("META-INF/MANIFEST.MF")); + mf.write(jos); + jos.closeEntry(); + } + var zip = tmp.resolve("no-impl.zip"); + try (var zos = new ZipOutputStream(new FileOutputStream(zip.toFile()))) { + zos.putNextEntry(new ZipEntry("my-ext-1.0.0-api.jar")); + Files.copy(apiJar, zos); + zos.closeEntry(); + } + var results = new ArtifactLayoutValidator().validate(zip); + assertTrue(results.stream().anyMatch(r -> r.getRuleCode().equals("BTRACE-LO-004") && r.isError())); + } + + @Test void nonExistentPathIsError() { + var results = new ArtifactLayoutValidator().validate(tmp.resolve("nonexistent.zip")); + assertTrue(results.stream().anyMatch(r -> r.getRuleCode().equals("BTRACE-LO-001") && r.isError())); + } +} +``` + +- [ ] **Step 2: Run test — expect compilation failure** + +```bash +./gradlew :btrace-ext-validator:test 2>&1 | grep "error:" +``` + +- [ ] **Step 3: Implement ArtifactLayoutValidator** + +Create `btrace-ext-validator/src/main/java/io/btrace/ext/validator/ArtifactLayoutValidator.java`: +```java +package io.btrace.ext.validator; + +import java.io.*; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.*; +import java.util.jar.*; +import java.util.zip.*; + +public final class ArtifactLayoutValidator { + + public List validate(Path extensionZip) { + var results = new ArrayList(); + if (!Files.exists(extensionZip)) { + results.add(new ValidationResult("BTRACE-LO-001", ValidationSeverity.ERROR, + "Extension ZIP not found: " + extensionZip, extensionZip.getFileName().toString())); + return results; + } + try (var zf = new ZipFile(extensionZip.toFile())) { + var entries = Collections.list(zf.entries()).stream() + .map(ZipEntry::getName).toList(); + + String apiEntry = entries.stream() + .filter(e -> e.endsWith("-api.jar")).findFirst().orElse(null); + if (apiEntry == null) { + results.add(new ValidationResult("BTRACE-LO-002", ValidationSeverity.ERROR, + "No *-api.jar found in extension ZIP", extensionZip.getFileName().toString())); + return results; + } + + Manifest mf = readManifest(zf, apiEntry, results); + if (mf == null) return results; + + String implName = mf.getMainAttributes().getValue("BTrace-Extension-Impl"); + if (implName == null || implName.isBlank()) { + // ManifestValidator will catch this — just skip impl check here + return results; + } + if (entries.stream().noneMatch(e -> e.equals(implName))) { + results.add(new ValidationResult("BTRACE-LO-004", ValidationSeverity.ERROR, + "Impl JAR '" + implName + "' declared in manifest not found in ZIP", + extensionZip.getFileName().toString())); + } + if (!implName.endsWith("-impl.jar")) { + results.add(new ValidationResult("BTRACE-LO-005", ValidationSeverity.WARN, + "Impl JAR name should follow '{name}-{version}-impl.jar' convention: " + implName, + implName)); + } + } catch (IOException e) { + results.add(new ValidationResult("BTRACE-LO-001", ValidationSeverity.ERROR, + "Cannot open extension ZIP: " + e.getMessage(), extensionZip.getFileName().toString())); + } + return results; + } + + private Manifest readManifest(ZipFile zf, String apiEntry, List results) { + var ze = zf.getEntry(apiEntry); + try (var apiIs = zf.getInputStream(ze); + var jis = new JarInputStream(apiIs)) { + Manifest mf = jis.getManifest(); + if (mf == null) { + results.add(new ValidationResult("BTRACE-LO-003", ValidationSeverity.ERROR, + "api.jar has no MANIFEST.MF", apiEntry)); + return null; + } + return mf; + } catch (IOException e) { + results.add(new ValidationResult("BTRACE-LO-003", ValidationSeverity.ERROR, + "Cannot read api.jar manifest: " + e.getMessage(), apiEntry)); + return null; + } + } +} +``` + +- [ ] **Step 4: Run tests — expect PASS** + +```bash +./gradlew :btrace-ext-validator:test +``` +Expected: `BUILD SUCCESSFUL`, all tests passed. + +- [ ] **Step 5: Commit** + +```bash +git add btrace-ext-validator/ +git commit -m "feat(tck): add ArtifactLayoutValidator with BTRACE-LO-* rules" +``` + +--- + +## Task 4: ServiceApiValidator + +**Files:** +- Create: `btrace-ext-validator/src/main/java/io/btrace/ext/validator/ServiceApiValidator.java` +- Create: `btrace-ext-validator/src/test/java/io/btrace/ext/validator/ServiceApiValidatorTest.java` + +Implements BTRACE-EXT-001, 002, 003, 010, 013, 020, 021, 022, 041 via pure ASM (no classloading). + +Forbidden type prefixes for BTRACE-EXT-013: +`java/io/`, `java/net/`, `java/nio/channels/`, `java/lang/reflect/` + +Nullable annotation descriptors (default): +`Ljavax/annotation/Nullable;`, `Lorg/jspecify/annotations/Nullable;`, `Lorg/jetbrains/annotations/Nullable;`, `Ljakarta/annotation/Nullable;` + +Nonnull annotation descriptors (default): +`Ljavax/annotation/Nonnull;`, `Lorg/jspecify/annotations/NonNull;`, `Lorg/jetbrains/annotations/NotNull;`, `Ljakarta/annotation/Nonnull;` + +- [ ] **Step 1: Write failing tests** + +Create `btrace-ext-validator/src/test/java/io/btrace/ext/validator/ServiceApiValidatorTest.java`: +```java +package io.btrace.ext.validator; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import org.objectweb.asm.*; +import org.objectweb.asm.Opcodes; +import java.io.*; +import java.nio.file.*; +import java.util.*; +import java.util.jar.*; +import java.util.zip.*; +import static org.junit.jupiter.api.Assertions.*; + +class ServiceApiValidatorTest { + @TempDir Path tmp; + + /** Write a class file to a temporary api.jar and return the jar path. */ + private Path apiJarWith(String className, byte[] classBytes) throws Exception { + var jar = tmp.resolve("test-api.jar"); + try (var jos = new JarOutputStream(new FileOutputStream(jar.toFile()))) { + var mf = new Manifest(); + mf.getMainAttributes().put(Attributes.Name.MANIFEST_VERSION, "1.0"); + mf.getMainAttributes().putValue("BTrace-Extension-Services", + className.replace('/', '.')); + jos.putNextEntry(new JarEntry("META-INF/MANIFEST.MF")); + mf.write(jos); + jos.closeEntry(); + jos.putNextEntry(new JarEntry(className + ".class")); + jos.write(classBytes); + jos.closeEntry(); + } + return jar; + } + + private byte[] publicInterface(String name) { + var cw = new ClassWriter(0); + cw.visit(Opcodes.V11, Opcodes.ACC_PUBLIC | Opcodes.ACC_INTERFACE | Opcodes.ACC_ABSTRACT, + name, null, "java/lang/Object", null); + // add one method with @Nullable return + var mv = cw.visitMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_ABSTRACT, + "doWork", "()Ljava/lang/String;", null, null); + mv.visitAnnotation("Ljavax/annotation/Nullable;", true).visitEnd(); + mv.visitEnd(); + cw.visitEnd(); + return cw.toByteArray(); + } + + private byte[] classNotInterface(String name) { + var cw = new ClassWriter(0); + cw.visit(Opcodes.V11, Opcodes.ACC_PUBLIC, name, null, "java/lang/Object", null); + cw.visitEnd(); + return cw.toByteArray(); + } + + private byte[] interfaceWithForbiddenParam(String name) { + var cw = new ClassWriter(0); + cw.visit(Opcodes.V11, Opcodes.ACC_PUBLIC | Opcodes.ACC_INTERFACE | Opcodes.ACC_ABSTRACT, + name, null, "java/lang/Object", null); + var mv = cw.visitMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_ABSTRACT, + "read", "(Ljava/io/InputStream;)V", null, null); + mv.visitEnd(); + cw.visitEnd(); + return cw.toByteArray(); + } + + @Test void validInterfaceProducesNoErrors() throws Exception { + var jar = apiJarWith("io/example/MyService", publicInterface("io/example/MyService")); + var results = new ServiceApiValidator().validate(jar); + assertTrue(results.stream().noneMatch(ValidationResult::isError), + () -> "Unexpected errors: " + results); + } + + @Test void classInsteadOfInterfaceIsExt001() throws Exception { + var jar = apiJarWith("io/example/MyService", classNotInterface("io/example/MyService")); + var results = new ServiceApiValidator().validate(jar); + assertTrue(results.stream().anyMatch(r -> r.getRuleCode().equals("BTRACE-EXT-001") && r.isError())); + } + + @Test void forbiddenParamTypeIsExt013() throws Exception { + var jar = apiJarWith("io/example/MyService", interfaceWithForbiddenParam("io/example/MyService")); + var results = new ServiceApiValidator().validate(jar); + assertTrue(results.stream().anyMatch(r -> r.getRuleCode().equals("BTRACE-EXT-013") && r.isError())); + } + + @Test void serviceClassMissingFromJarIsExt041() throws Exception { + // Manifest declares a service that isn't in the JAR + var jar = tmp.resolve("missing-service.jar"); + try (var jos = new JarOutputStream(new FileOutputStream(jar.toFile()))) { + var mf = new Manifest(); + mf.getMainAttributes().put(Attributes.Name.MANIFEST_VERSION, "1.0"); + mf.getMainAttributes().putValue("BTrace-Extension-Services", "io.example.Ghost"); + jos.putNextEntry(new JarEntry("META-INF/MANIFEST.MF")); + mf.write(jos); + jos.closeEntry(); + } + var results = new ServiceApiValidator().validate(jar); + assertTrue(results.stream().anyMatch(r -> r.getRuleCode().equals("BTRACE-EXT-041") && r.isError())); + } +} +``` + +- [ ] **Step 2: Run test — expect compilation failure** + +```bash +./gradlew :btrace-ext-validator:test 2>&1 | grep "error:" +``` + +- [ ] **Step 3: Implement ServiceApiValidator** + +Create `btrace-ext-validator/src/main/java/io/btrace/ext/validator/ServiceApiValidator.java`: +```java +package io.btrace.ext.validator; + +import org.objectweb.asm.*; +import org.objectweb.asm.Opcodes; +import org.objectweb.asm.tree.*; +import java.io.*; +import java.nio.file.Path; +import java.util.*; +import java.util.jar.*; + +public final class ServiceApiValidator { + + private static final Set FORBIDDEN_PREFIXES = Set.of( + "java/io/", "java/net/", "java/nio/channels/", "java/lang/reflect/"); + + private static final Set NULLABLE_DESCS = Set.of( + "Ljavax/annotation/Nullable;", "Lorg/jspecify/annotations/Nullable;", + "Lorg/jetbrains/annotations/Nullable;", "Ljakarta/annotation/Nullable;"); + + private static final Set NONNULL_DESCS = Set.of( + "Ljavax/annotation/Nonnull;", "Lorg/jspecify/annotations/NonNull;", + "Lorg/jetbrains/annotations/NotNull;", "Ljakarta/annotation/Nonnull;"); + + public List validate(Path apiJar) { + var results = new ArrayList(); + try (var jf = new JarFile(apiJar.toFile())) { + var mf = jf.getManifest(); + if (mf == null) { + results.add(err("BTRACE-EXT-041", "api.jar has no MANIFEST.MF", apiJar.getFileName().toString())); + return results; + } + String svcAttr = mf.getMainAttributes().getValue("BTrace-Extension-Services"); + if (svcAttr == null || svcAttr.isBlank()) return results; + + var declaredServices = Arrays.stream(svcAttr.split(",")) + .map(String::trim).filter(s -> !s.isEmpty()) + .map(s -> s.replace('.', '/')).toList(); + + var classEntries = new HashMap(); + var entries = jf.entries(); + while (entries.hasMoreElements()) { + var entry = entries.nextElement(); + if (entry.getName().endsWith(".class") && !entry.isDirectory()) { + classEntries.put( + entry.getName().replace('/', '.').replaceAll("\\.class$", ""), + jf.getInputStream(entry).readAllBytes()); + } + } + + for (String svcInternal : declaredServices) { + String svcFqcn = svcInternal.replace('/', '.'); + byte[] classBytes = classEntries.get(svcFqcn); + if (classBytes == null) { + results.add(err("BTRACE-EXT-041", + "Declared service '" + svcFqcn + "' not found in api.jar", apiJar.getFileName().toString())); + continue; + } + validateServiceClass(svcFqcn, classBytes, results); + } + } catch (IOException e) { + results.add(err("BTRACE-EXT-041", "Cannot read api.jar: " + e.getMessage(), apiJar.getFileName().toString())); + } + return results; + } + + private void validateServiceClass(String fqcn, byte[] bytes, List results) { + var cr = new ClassReader(bytes); + var cn = new ClassNode(); + cr.accept(cn, ClassReader.SKIP_CODE | ClassReader.SKIP_DEBUG); + + boolean isInterface = (cn.access & Opcodes.ACC_INTERFACE) != 0; + boolean isPublic = (cn.access & Opcodes.ACC_PUBLIC) != 0; + boolean isTopLevel = cn.outerClass == null; + + if (!isInterface || !isPublic || !isTopLevel) { + results.add(err("BTRACE-EXT-001", + "Service '" + fqcn + "' must be a public, top-level interface", fqcn)); + } + + for (FieldNode f : cn.fields) { + boolean isConst = (f.access & (Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC | Opcodes.ACC_FINAL)) + == (Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC | Opcodes.ACC_FINAL); + if (!isConst) { + results.add(err("BTRACE-EXT-003", + "Non-constant field '" + f.name + "' in service '" + fqcn + "'", fqcn)); + } + } + + for (MethodNode mn : cn.methods) { + if ("".equals(mn.name)) continue; + boolean isDefault = !((mn.access & Opcodes.ACC_ABSTRACT) != 0); + boolean isPrivate = (mn.access & Opcodes.ACC_PRIVATE) != 0; + if (isDefault || isPrivate) { + results.add(err("BTRACE-EXT-002", + "Method '" + mn.name + "' in '" + fqcn + "' must not be default or private", fqcn)); + } + checkExceptions(fqcn, mn, results); + checkSignatureForForbiddenTypes(fqcn, mn, results); + checkNullability(fqcn, mn, results); + } + } + + private void checkExceptions(String fqcn, MethodNode mn, List results) { + if (mn.exceptions == null) return; + for (String ex : mn.exceptions) { + if (!isUnchecked(ex)) { + results.add(err("BTRACE-EXT-010", + "Method '" + mn.name + "' in '" + fqcn + "' throws checked exception: " + ex.replace('/', '.'), fqcn)); + } + } + } + + private boolean isUnchecked(String internalName) { + return internalName.startsWith("java/lang/RuntimeException") + || internalName.startsWith("java/lang/Error") + || internalName.equals("java/lang/Throwable"); + } + + private void checkSignatureForForbiddenTypes(String fqcn, MethodNode mn, List results) { + var desc = mn.desc; + for (Type t : getAllTypes(desc)) { + if (t.getSort() == Type.OBJECT) { + String internalName = t.getInternalName(); + for (String prefix : FORBIDDEN_PREFIXES) { + if (internalName.startsWith(prefix)) { + results.add(err("BTRACE-EXT-013", + "Method '" + mn.name + "' in '" + fqcn + "' uses forbidden type: " + + internalName.replace('/', '.'), fqcn)); + } + } + } + } + } + + private List getAllTypes(String methodDesc) { + var types = new ArrayList(); + try { + types.addAll(Arrays.asList(Type.getArgumentTypes(methodDesc))); + Type ret = Type.getReturnType(methodDesc); + if (ret.getSort() != Type.VOID) types.add(ret); + } catch (Exception ignored) {} + return types; + } + + private void checkNullability(String fqcn, MethodNode mn, List results) { + Type returnType = Type.getReturnType(mn.desc); + boolean voidReturn = returnType.getSort() == Type.VOID; + boolean primitiveReturn = returnType.getSort() != Type.OBJECT && returnType.getSort() != Type.ARRAY; + + if (!voidReturn && !primitiveReturn) { + boolean annotated = hasAnnotation(mn.visibleAnnotations, NULLABLE_DESCS) + || hasAnnotation(mn.visibleAnnotations, NONNULL_DESCS) + || hasAnnotation(mn.invisibleAnnotations, NULLABLE_DESCS) + || hasAnnotation(mn.invisibleAnnotations, NONNULL_DESCS); + if (!annotated) { + results.add(warn("BTRACE-EXT-020", + "Return type of '" + mn.name + "' in '" + fqcn + "' lacks @Nullable/@Nonnull", fqcn)); + } + if (returnType.getSort() == Type.OBJECT) { + boolean isInterface = returnType.getInternalName().contains("/"); + boolean isNullable = hasAnnotation(mn.visibleAnnotations, NULLABLE_DESCS) + || hasAnnotation(mn.invisibleAnnotations, NULLABLE_DESCS); + if (!isNullable) { + results.add(err("BTRACE-EXT-022", + "Interface return type of '" + mn.name + "' in '" + fqcn + "' must be @Nullable", fqcn)); + } + } + } + + Type[] argTypes = Type.getArgumentTypes(mn.desc); + List[] paramAnnotations = mn.visibleParameterAnnotations; + for (int i = 0; i < argTypes.length; i++) { + if (argTypes[i].getSort() == Type.OBJECT || argTypes[i].getSort() == Type.ARRAY) { + List annots = (paramAnnotations != null && i < paramAnnotations.length) + ? paramAnnotations[i] : null; + boolean annotated = hasAnnotation(annots, NULLABLE_DESCS) + || hasAnnotation(annots, NONNULL_DESCS); + if (!annotated) { + results.add(warn("BTRACE-EXT-021", + "Parameter " + i + " of '" + mn.name + "' in '" + fqcn + "' lacks @Nullable/@Nonnull", fqcn)); + } + } + } + } + + private boolean hasAnnotation(List annots, Set descs) { + if (annots == null) return false; + return annots.stream().anyMatch(a -> descs.contains(a.desc)); + } + + private static ValidationResult err(String code, String msg, String artifact) { + return new ValidationResult(code, ValidationSeverity.ERROR, msg, artifact); + } + + private static ValidationResult warn(String code, String msg, String artifact) { + return new ValidationResult(code, ValidationSeverity.WARN, msg, artifact); + } +} +``` + +- [ ] **Step 4: Run tests — expect PASS** + +```bash +./gradlew :btrace-ext-validator:test +``` +Expected: `BUILD SUCCESSFUL`. + +- [ ] **Step 5: Commit** + +```bash +git add btrace-ext-validator/ +git commit -m "feat(tck): add ServiceApiValidator with BTRACE-EXT-* rules" +``` + +--- + +## Task 5: ApiImplPartitionValidator + +**Files:** +- Create: `btrace-ext-validator/src/main/java/io/btrace/ext/validator/ApiImplPartitionValidator.java` +- Create: `btrace-ext-validator/src/test/java/io/btrace/ext/validator/ApiImplPartitionValidatorTest.java` + +Rule codes: +- `BTRACE-PT-001` ERROR — impl-only class appears in an api.jar method signature + +- [ ] **Step 1: Write failing test** + +Create `btrace-ext-validator/src/test/java/io/btrace/ext/validator/ApiImplPartitionValidatorTest.java`: +```java +package io.btrace.ext.validator; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import org.objectweb.asm.*; +import org.objectweb.asm.Opcodes; +import java.io.*; +import java.nio.file.*; +import java.util.*; +import java.util.jar.*; +import static org.junit.jupiter.api.Assertions.*; + +class ApiImplPartitionValidatorTest { + @TempDir Path tmp; + + private Path buildApiJar(String serviceClass, byte[] serviceBytes, + String... extraClasses) throws Exception { + var jar = tmp.resolve("api.jar"); + try (var jos = new JarOutputStream(new FileOutputStream(jar.toFile()))) { + var mf = new Manifest(); + mf.getMainAttributes().put(Attributes.Name.MANIFEST_VERSION, "1.0"); + mf.getMainAttributes().putValue("BTrace-Extension-Services", serviceClass.replace('/', '.')); + jos.putNextEntry(new JarEntry("META-INF/MANIFEST.MF")); + mf.write(jos); + jos.closeEntry(); + jos.putNextEntry(new JarEntry(serviceClass + ".class")); + jos.write(serviceBytes); + jos.closeEntry(); + } + return jar; + } + + private Path buildImplJar(String implClass) throws Exception { + var jar = tmp.resolve("impl.jar"); + try (var jos = new JarOutputStream(new FileOutputStream(jar.toFile()))) { + var cw = new ClassWriter(0); + cw.visit(Opcodes.V11, Opcodes.ACC_PUBLIC, implClass, null, "java/lang/Object", null); + cw.visitEnd(); + jos.putNextEntry(new JarEntry(implClass + ".class")); + jos.write(cw.toByteArray()); + jos.closeEntry(); + } + return jar; + } + + @Test void cleanPartitionProducesNoErrors() throws Exception { + var cw = new ClassWriter(0); + cw.visit(Opcodes.V11, Opcodes.ACC_PUBLIC | Opcodes.ACC_INTERFACE | Opcodes.ACC_ABSTRACT, + "io/example/MyService", null, "java/lang/Object", null); + cw.visitMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_ABSTRACT, "doWork", "()V", null, null).visitEnd(); + cw.visitEnd(); + + var api = buildApiJar("io/example/MyService", cw.toByteArray()); + var impl = buildImplJar("io/example/internal/MyServiceImpl"); + var results = new ApiImplPartitionValidator().validate(api, impl); + assertTrue(results.stream().noneMatch(ValidationResult::isError)); + } + + @Test void implTypeInApiSignatureIsError() throws Exception { + var cw = new ClassWriter(0); + cw.visit(Opcodes.V11, Opcodes.ACC_PUBLIC | Opcodes.ACC_INTERFACE | Opcodes.ACC_ABSTRACT, + "io/example/MyService", null, "java/lang/Object", null); + // Method that returns an impl-only type + cw.visitMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_ABSTRACT, + "getImpl", "()Lio/example/internal/MyServiceImpl;", null, null).visitEnd(); + cw.visitEnd(); + + var api = buildApiJar("io/example/MyService", cw.toByteArray()); + var impl = buildImplJar("io/example/internal/MyServiceImpl"); + var results = new ApiImplPartitionValidator().validate(api, impl); + assertTrue(results.stream().anyMatch(r -> r.getRuleCode().equals("BTRACE-PT-001") && r.isError())); + } +} +``` + +- [ ] **Step 2: Implement ApiImplPartitionValidator** + +Create `btrace-ext-validator/src/main/java/io/btrace/ext/validator/ApiImplPartitionValidator.java`: +```java +package io.btrace.ext.validator; + +import org.objectweb.asm.*; +import org.objectweb.asm.Opcodes; +import org.objectweb.asm.tree.*; +import java.io.*; +import java.nio.file.Path; +import java.util.*; +import java.util.jar.*; + +public final class ApiImplPartitionValidator { + + public List validate(Path apiJar, Path implJar) { + var results = new ArrayList(); + var implTypes = collectClassNames(implJar, results); + if (!results.isEmpty()) return results; + + try (var jf = new JarFile(apiJar.toFile())) { + var entries = jf.entries(); + while (entries.hasMoreElements()) { + var entry = entries.nextElement(); + if (!entry.getName().endsWith(".class") || entry.isDirectory()) continue; + byte[] bytes = jf.getInputStream(entry).readAllBytes(); + var cr = new ClassReader(bytes); + var cn = new ClassNode(); + cr.accept(cn, ClassReader.SKIP_CODE | ClassReader.SKIP_DEBUG); + String apiClass = cn.name.replace('/', '.'); + for (MethodNode mn : cn.methods) { + checkDesc(apiClass, mn.name, mn.desc, implTypes, results); + } + } + } catch (IOException e) { + results.add(new ValidationResult("BTRACE-PT-001", ValidationSeverity.ERROR, + "Cannot read api.jar: " + e.getMessage(), apiJar.getFileName().toString())); + } + return results; + } + + private Set collectClassNames(Path jar, List errors) { + var names = new HashSet(); + try (var jf = new JarFile(jar.toFile())) { + var entries = jf.entries(); + while (entries.hasMoreElements()) { + var entry = entries.nextElement(); + if (entry.getName().endsWith(".class") && !entry.isDirectory()) { + names.add(entry.getName().replace('/', '.').replaceAll("\\.class$", "")); + } + } + } catch (IOException e) { + errors.add(new ValidationResult("BTRACE-PT-001", ValidationSeverity.ERROR, + "Cannot read impl.jar: " + e.getMessage(), jar.getFileName().toString())); + } + return names; + } + + private void checkDesc(String apiClass, String method, String desc, + Set implTypes, List results) { + for (Type t : allTypes(desc)) { + if (t.getSort() == Type.OBJECT) { + String fqcn = t.getClassName(); + if (implTypes.contains(fqcn)) { + results.add(new ValidationResult("BTRACE-PT-001", ValidationSeverity.ERROR, + "Impl-only type '" + fqcn + "' appears in API method '" + + apiClass + "." + method + "'", apiClass)); + } + } + } + } + + private List allTypes(String desc) { + var out = new ArrayList(); + try { + out.addAll(Arrays.asList(Type.getArgumentTypes(desc))); + var ret = Type.getReturnType(desc); + if (ret.getSort() != Type.VOID) out.add(ret); + } catch (Exception ignored) {} + return out; + } +} +``` + +- [ ] **Step 3: Run tests — expect PASS** + +```bash +./gradlew :btrace-ext-validator:test +``` + +- [ ] **Step 4: Commit** + +```bash +git add btrace-ext-validator/ +git commit -m "feat(tck): add ApiImplPartitionValidator with BTRACE-PT-001" +``` + +--- + +## Task 6: Wire btrace-ext-validator into Gradle plugin + +**Files:** +- Modify: `btrace-gradle-plugin/build.gradle` +- Modify: `btrace-gradle-plugin/src/main/groovy/io/btrace/gradle/BTraceExtensionPlugin.groovy` + +The goal: replace the inline ASM bytecode analysis in `validateServiceApis` with a call to `ServiceApiValidator`. The Groovy task body shrinks to: iterate class dirs, call validator, translate `ValidationResult` list to Gradle warnings/errors. + +- [ ] **Step 1: Add dependency in plugin build.gradle** + +In `btrace-gradle-plugin/build.gradle`, add to `dependencies`: +```groovy +implementation project(':btrace-ext-validator') +``` + +Also add to `settings.gradle` (the plugin included build needs the module on the classpath — since `btrace-gradle-plugin` is an `includeBuild`, add the dependency resolution): +```groovy +// in btrace-gradle-plugin/settings.gradle (create if absent) +includeBuild '..' // gives access to root project's modules +``` + +> **Note:** The included build resolution for `project(':btrace-ext-validator')` may require using `files(...)` pointing to the built JAR if Gradle's composite build doesn't resolve cross-project dependencies here. Check by running `./gradlew :btrace-gradle-plugin:dependencies` and adjust if needed — substitution with `implementation(files('../btrace-ext-validator/build/libs/btrace-ext-validator.jar'))` is acceptable as a fallback. + +- [ ] **Step 2: Refactor validateServiceApis task** + +In `BTraceExtensionPlugin.groovy`, find the `validateServiceApis` task definition (around line 196) and replace its `doLast` body with: + +```groovy +def validateServiceApis = project.tasks.register('validateServiceApis') { + dependsOn(authoredCompileTask) + doLast { + def apiJarFile = buildApiJar.get().archiveFile.get().asFile.toPath() + if (!apiJarFile.toFile().exists()) { + project.logger.warn('[BTRACE-EXT] api.jar not built yet; skipping validation') + return + } + def validator = new io.btrace.ext.validator.ServiceApiValidator() + def results = validator.validate(apiJarFile) + def sev = (project.extensions.findByType(BTraceExtensionMetadata)?.nullabilitySeverity ?: 'warn') + results.each { r -> + def line = "[${r.ruleCode}] ${r.message} (${r.artifact})" + if (r.isError() || "error".equalsIgnoreCase(sev)) { + throw new GradleException("[BTRACE-EXT] Validation failed: $line") + } else { + project.logger.warn("[BTRACE-EXT] $line") + } + } + } +} +``` + +- [ ] **Step 3: Verify plugin tests still pass** + +```bash +./gradlew :btrace-gradle-plugin:test +``` +Expected: `BUILD SUCCESSFUL`. + +- [ ] **Step 4: Run full check on btrace-contracts to verify end-to-end** + +```bash +./gradlew :btrace-extensions:btrace-contracts:check +``` +Expected: `BUILD SUCCESSFUL` with no new warnings. + +- [ ] **Step 5: Commit** + +```bash +git add btrace-gradle-plugin/ +git commit -m "refactor(tck): delegate validateServiceApis to ServiceApiValidator" +``` + +--- + +## Task 7: btrace-tck module skeleton + +**Files:** +- Create: `btrace-tck/build.gradle` +- Create: `btrace-tck/src/main/java/io/btrace/tck/TckStatus.java` +- Create: `btrace-tck/src/main/java/io/btrace/tck/TckResult.java` +- Create: `btrace-tck/src/main/java/io/btrace/tck/TckSuiteResult.java` +- Create: `btrace-tck/src/main/java/io/btrace/tck/TckInput.java` +- Create: `btrace-tck/src/main/java/io/btrace/tck/TckEngine.java` +- Modify: `settings.gradle` + +- [ ] **Step 1: Add to settings.gradle** + +```groovy +include 'btrace-tck' +``` + +- [ ] **Step 2: Create btrace-tck/build.gradle** + +```groovy +apply from: rootProject.file('common.gradle') + +java { + toolchain { languageVersion = JavaLanguageVersion.of(11) } +} + +jar { + manifest { attributes 'Main-Class': 'io.btrace.tck.cli.TckMain' } + from { configurations.runtimeClasspath.collect { it.isDirectory() ? it : zipTree(it) } } + duplicatesStrategy = 'exclude' +} + +dependencies { + implementation project(':btrace-ext-validator') + implementation 'com.fasterxml.jackson.core:jackson-databind:2.18.2' + implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml:2.18.2' + implementation 'info.picocli:picocli:4.7.6' + testImplementation platform('org.junit:junit-bom:5.11.4') + testImplementation 'org.junit.jupiter:junit-jupiter' + testRuntimeOnly 'org.junit.platform:junit-platform-launcher' +} + +test { useJUnitPlatform() } +``` + +- [ ] **Step 3: Write failing test** + +Create `btrace-tck/src/test/java/io/btrace/tck/TckEngineTest.java`: +```java +package io.btrace.tck; + +import org.junit.jupiter.api.Test; +import java.util.List; +import static org.junit.jupiter.api.Assertions.*; + +class TckEngineTest { + @Test void engineWithNoSuitesReturnsEmptyResults() { + var engine = new TckEngine(List.of()); + var input = TckInput.builder().build(); + var results = engine.run(input); + assertTrue(results.isEmpty()); + } + + @Test void skipsDownstreamSuiteOnUpstreamFailure() { + var failSuite = new StubSuite("fail", TckStatus.FAIL); + var skipSuite = new StubSuite("skip", TckStatus.PASS); + var engine = new TckEngine(List.of(failSuite, skipSuite)); + var results = engine.run(TckInput.builder().build()); + assertEquals(2, results.size()); + assertEquals(TckStatus.FAIL, results.get(0).overallStatus()); + assertEquals(TckStatus.SKIP, results.get(1).overallStatus()); + } + + private record StubSuite(String name, TckStatus status) implements Suite { + public TckSuiteResult run(TckInput input) { + return new TckSuiteResult(name, List.of( + new TckResult(name, "check", status, null, null))); + } + public String name() { return name; } + } +} +``` + +- [ ] **Step 4: Create core TCK types** + +`TckStatus.java`: +```java +package io.btrace.tck; +public enum TckStatus { PASS, FAIL, SKIP } +``` + +`TckResult.java`: +```java +package io.btrace.tck; + +public record TckResult( + String suiteName, + String checkName, + TckStatus status, + String ruleCode, // nullable — present on FAIL + String message // nullable +) { + public boolean isFail() { return status == TckStatus.FAIL; } +} +``` + +`TckSuiteResult.java`: +```java +package io.btrace.tck; + +import java.util.List; + +public record TckSuiteResult(String suiteName, List checks) { + public TckStatus overallStatus() { + if (checks.isEmpty()) return TckStatus.PASS; + if (checks.stream().anyMatch(r -> r.status() == TckStatus.SKIP)) return TckStatus.SKIP; + return checks.stream().anyMatch(TckResult::isFail) ? TckStatus.FAIL : TckStatus.PASS; + } + public boolean hasFailed() { return overallStatus() == TckStatus.FAIL; } +} +``` + +`TckInput.java`: +```java +package io.btrace.tck; + +import java.nio.file.Path; + +public final class TckInput { + private final Path extensionZip; + private final Path btraceHome; + private final Path tckConfig; + private final Path reportDir; + + private TckInput(Builder b) { + this.extensionZip = b.extensionZip; + this.btraceHome = b.btraceHome; + this.tckConfig = b.tckConfig; + this.reportDir = b.reportDir; + } + + public Path extensionZip() { return extensionZip; } + public Path btraceHome() { return btraceHome; } + public Path tckConfig() { return tckConfig; } + public Path reportDir() { return reportDir; } + + public static Builder builder() { return new Builder(); } + + public static final class Builder { + private Path extensionZip; + private Path btraceHome; + private Path tckConfig; + private Path reportDir = Path.of("tck-report"); + + public Builder extensionZip(Path p) { this.extensionZip = p; return this; } + public Builder btraceHome(Path p) { this.btraceHome = p; return this; } + public Builder tckConfig(Path p) { this.tckConfig = p; return this; } + public Builder reportDir(Path p) { this.reportDir = p; return this; } + public TckInput build() { return new TckInput(this); } + } +} +``` + +`Suite.java` (interface): +```java +package io.btrace.tck; +public interface Suite { + String name(); + TckSuiteResult run(TckInput input); +} +``` + +`TckEngine.java`: +```java +package io.btrace.tck; + +import java.util.ArrayList; +import java.util.List; + +public final class TckEngine { + private final List suites; + + public TckEngine(List suites) { + this.suites = List.copyOf(suites); + } + + public List run(TckInput input) { + var results = new ArrayList(); + boolean failed = false; + for (Suite suite : suites) { + if (failed) { + results.add(skipSuite(suite.name())); + } else { + var result = suite.run(input); + results.add(result); + if (result.hasFailed()) failed = true; + } + } + return results; + } + + private TckSuiteResult skipSuite(String name) { + return new TckSuiteResult(name, List.of( + new TckResult(name, "*", TckStatus.SKIP, null, + "Upstream suite failed — skipped"))); + } +} +``` + +- [ ] **Step 5: Run tests — expect PASS** + +```bash +./gradlew :btrace-tck:test +``` +Expected: `BUILD SUCCESSFUL`. + +- [ ] **Step 6: Commit** + +```bash +git add btrace-tck/ settings.gradle +git commit -m "feat(tck): add btrace-tck module skeleton with TckEngine" +``` + +--- + +## Task 8: Reporting infrastructure + +**Files:** +- Create: `btrace-tck/src/main/java/io/btrace/tck/report/StdoutReporter.java` +- Create: `btrace-tck/src/main/java/io/btrace/tck/report/JUnitXmlReporter.java` +- Create: `btrace-tck/src/main/java/io/btrace/tck/report/HtmlReporter.java` +- Create: `btrace-tck/src/test/java/io/btrace/tck/report/ReporterTest.java` + +- [ ] **Step 1: Write failing tests** + +Create `btrace-tck/src/test/java/io/btrace/tck/report/ReporterTest.java`: +```java +package io.btrace.tck.report; + +import io.btrace.tck.*; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import java.io.*; +import java.nio.file.*; +import java.util.List; +import static org.junit.jupiter.api.Assertions.*; + +class ReporterTest { + private List sampleResults() { + return List.of( + new TckSuiteResult("structural", List.of( + new TckResult("structural", "ManifestCheck", TckStatus.PASS, null, null), + new TckResult("structural", "ServiceApiCheck", TckStatus.FAIL, + "BTRACE-EXT-001", "Service must be interface"))), + new TckSuiteResult("lifecycle", List.of( + new TckResult("lifecycle", "*", TckStatus.SKIP, null, "Upstream failed")))); + } + + @Test void stdoutReporterContainsPassFailSkip() { + var sw = new StringWriter(); + new StdoutReporter(new PrintWriter(sw)).report(sampleResults()); + var out = sw.toString(); + assertTrue(out.contains("[PASS]")); + assertTrue(out.contains("[FAIL]")); + assertTrue(out.contains("[SKIP]")); + assertTrue(out.contains("BTRACE-EXT-001")); + } + + @Test void junitXmlContainsTestcases(@TempDir Path tmp) throws Exception { + var out = tmp.resolve("results.xml"); + new JUnitXmlReporter(out).report(sampleResults()); + var xml = Files.readString(out); + assertTrue(xml.contains(" suites) { + long fails = 0, skips = 0; + for (var suite : suites) { + for (var check : suite.checks()) { + String tag = switch (check.status()) { + case PASS -> "[PASS]"; + case FAIL -> "[FAIL]"; + case SKIP -> "[SKIP]"; + }; + String msg = suite.suiteName() + "/" + check.checkName(); + if (check.ruleCode() != null) msg += " " + check.ruleCode(); + if (check.message() != null) msg += " " + check.message(); + out.println(tag + " " + msg); + if (check.isFail()) fails++; + if (check.status() == TckStatus.SKIP) skips++; + } + } + out.println(); + boolean passed = fails == 0; + out.printf("TCK %s %d failure(s), %d suite(s) skipped%n", + passed ? "PASSED" : "FAILED", fails, skips); + } +} +``` + +- [ ] **Step 3: Implement JUnitXmlReporter** + +Create `btrace-tck/src/main/java/io/btrace/tck/report/JUnitXmlReporter.java`: +```java +package io.btrace.tck.report; + +import io.btrace.tck.*; +import java.io.*; +import java.nio.file.Path; +import java.util.List; + +public final class JUnitXmlReporter { + private final Path output; + + public JUnitXmlReporter(Path output) { this.output = output; } + + public void report(List suites) throws IOException { + output.getParent().toFile().mkdirs(); + try (var w = new PrintWriter(new FileWriter(output.toFile()))) { + w.println(""); + w.println(""); + for (var suite : suites) { + long failures = suite.checks().stream().filter(TckResult::isFail).count(); + long skipped = suite.checks().stream() + .filter(r -> r.status() == TckStatus.SKIP).count(); + w.printf(" %n", + esc(suite.suiteName()), suite.checks().size(), failures, skipped); + for (var check : suite.checks()) { + w.printf(" %n", + esc(check.checkName()), esc(suite.suiteName())); + if (check.isFail()) { + w.printf(" %n", + esc(check.ruleCode() != null ? check.ruleCode() : "FAIL"), + esc(check.message() != null ? check.message() : "")); + } else if (check.status() == TckStatus.SKIP) { + w.println(" "); + } + w.println(" "); + } + w.println(" "); + } + w.println(""); + } + } + + private String esc(String s) { + return s == null ? "" : s.replace("&","&").replace("<","<").replace("\"","""); + } +} +``` + +- [ ] **Step 4: Implement HtmlReporter** + +Create `btrace-tck/src/main/java/io/btrace/tck/report/HtmlReporter.java`: +```java +package io.btrace.tck.report; + +import io.btrace.tck.*; +import java.io.*; +import java.nio.file.Path; +import java.util.List; + +public final class HtmlReporter { + private final Path output; + + public HtmlReporter(Path output) { this.output = output; } + + public void report(List suites) throws IOException { + output.getParent().toFile().mkdirs(); + try (var w = new PrintWriter(new FileWriter(output.toFile()))) { + w.println(""" + + BTrace TCK Report + +

BTrace Extension TCK Report

+ """); + + long totalFail = suites.stream() + .flatMap(s -> s.checks().stream()).filter(TckResult::isFail).count(); + w.printf("

Overall: %s — %d failure(s)

%n", + totalFail == 0 ? "PASS" : "FAIL", + totalFail == 0 ? "PASS" : "FAIL", totalFail); + + for (var suite : suites) { + w.printf("

%s

%n", esc(suite.suiteName())); + w.println(""); + for (var check : suite.checks()) { + w.printf("%n", + esc(check.checkName()), + check.status().name(), + check.status().name(), + esc(check.ruleCode()), + esc(check.message())); + } + w.println("
CheckStatusRuleMessage
%s%s%s%s
"); + } + w.println(""); + } + } + + private String esc(String s) { + return s == null ? "" : s.replace("&","&").replace("<","<"); + } +} +``` + +- [ ] **Step 5: Run tests — expect PASS** + +```bash +./gradlew :btrace-tck:test +``` + +- [ ] **Step 6: Commit** + +```bash +git add btrace-tck/ +git commit -m "feat(tck): add StdoutReporter, JUnitXmlReporter, HtmlReporter" +``` + +--- + +## Task 9: StructuralSuite + +**Files:** +- Create: `btrace-tck/src/main/java/io/btrace/tck/suite/StructuralSuite.java` +- Create: `btrace-tck/src/main/java/io/btrace/tck/check/structural/ManifestCheck.java` +- Create: `btrace-tck/src/main/java/io/btrace/tck/check/structural/ArtifactLayoutCheck.java` +- Create: `btrace-tck/src/main/java/io/btrace/tck/check/structural/ServiceApiCheck.java` +- Create: `btrace-tck/src/main/java/io/btrace/tck/check/structural/ApiImplPartitionCheck.java` +- Create: `btrace-tck/src/test/java/io/btrace/tck/suite/StructuralSuiteTest.java` + +Each check follows the same pattern: extract relevant artifact from the extension ZIP, delegate to the validator, map `ValidationResult` to `TckResult`. + +- [ ] **Step 1: Write failing test for StructuralSuite** + +Create `btrace-tck/src/test/java/io/btrace/tck/suite/StructuralSuiteTest.java`: +```java +package io.btrace.tck.suite; + +import io.btrace.tck.*; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import java.io.*; +import java.nio.file.*; +import java.util.*; +import java.util.jar.*; +import java.util.zip.*; +import static org.junit.jupiter.api.Assertions.*; + +class StructuralSuiteTest { + @TempDir Path tmp; + + @Test void suiteIsNamedStructural() { + assertEquals("structural", new StructuralSuite().name()); + } + + @Test void missingZipProducesFailResult() { + var input = TckInput.builder() + .extensionZip(tmp.resolve("nonexistent.zip")) + .reportDir(tmp.resolve("report")) + .build(); + var result = new StructuralSuite().run(input); + assertTrue(result.hasFailed(), "Expected FAIL for missing ZIP"); + } + + @Test void nullZipPathProducesFailResult() { + var input = TckInput.builder().reportDir(tmp.resolve("report")).build(); + var result = new StructuralSuite().run(input); + assertTrue(result.hasFailed()); + } +} +``` + +- [ ] **Step 2: Create check helpers and StructuralSuite** + +`ManifestCheck.java`: +```java +package io.btrace.tck.check.structural; + +import io.btrace.ext.validator.*; +import io.btrace.tck.*; +import java.io.*; +import java.nio.file.Path; +import java.util.*; +import java.util.jar.*; +import java.util.zip.*; + +public final class ManifestCheck { + private static final String NAME = "ManifestCheck"; + + public TckResult run(Path extensionZip) { + try (var zf = new ZipFile(extensionZip.toFile())) { + String apiEntry = Collections.list(zf.entries()).stream() + .map(ZipEntry::getName) + .filter(n -> n.endsWith("-api.jar")) + .findFirst().orElse(null); + if (apiEntry == null) { + return fail("structural", NAME, "BTRACE-LO-002", "No *-api.jar in ZIP"); + } + var ze = zf.getEntry(apiEntry); + try (var jis = new JarInputStream(zf.getInputStream(ze))) { + Manifest mf = jis.getManifest(); + if (mf == null) return fail("structural", NAME, "BTRACE-LO-003", "api.jar has no MANIFEST.MF"); + var results = new ManifestValidator().validate(mf); + var first = results.stream().filter(ValidationResult::isError).findFirst(); + if (first.isPresent()) { + return fail("structural", NAME, first.get().getRuleCode(), first.get().getMessage()); + } + return pass("structural", NAME); + } + } catch (IOException e) { + return fail("structural", NAME, "BTRACE-LO-001", "Cannot open ZIP: " + e.getMessage()); + } + } + + private TckResult pass(String suite, String check) { + return new TckResult(suite, check, TckStatus.PASS, null, null); + } + private TckResult fail(String suite, String check, String code, String msg) { + return new TckResult(suite, check, TckStatus.FAIL, code, msg); + } +} +``` + +`ArtifactLayoutCheck.java`: +```java +package io.btrace.tck.check.structural; + +import io.btrace.ext.validator.*; +import io.btrace.tck.*; +import java.nio.file.Path; + +public final class ArtifactLayoutCheck { + private static final String NAME = "ArtifactLayoutCheck"; + + public TckResult run(Path extensionZip) { + var results = new ArtifactLayoutValidator().validate(extensionZip); + var first = results.stream().filter(ValidationResult::isError).findFirst(); + if (first.isPresent()) { + return new TckResult("structural", NAME, TckStatus.FAIL, + first.get().getRuleCode(), first.get().getMessage()); + } + return new TckResult("structural", NAME, TckStatus.PASS, null, null); + } +} +``` + +`ServiceApiCheck.java`: +```java +package io.btrace.tck.check.structural; + +import io.btrace.ext.validator.*; +import io.btrace.tck.*; +import java.io.*; +import java.nio.file.*; +import java.util.*; +import java.util.zip.*; + +public final class ServiceApiCheck { + private static final String NAME = "ServiceApiCheck"; + + public TckResult run(Path extensionZip) { + try { + Path apiJar = extractApiJar(extensionZip); + if (apiJar == null) + return new TckResult("structural", NAME, TckStatus.FAIL, "BTRACE-LO-002", "No api.jar found"); + var results = new ServiceApiValidator().validate(apiJar); + var first = results.stream().filter(ValidationResult::isError).findFirst(); + if (first.isPresent()) + return new TckResult("structural", NAME, TckStatus.FAIL, + first.get().getRuleCode(), first.get().getMessage()); + return new TckResult("structural", NAME, TckStatus.PASS, null, null); + } catch (IOException e) { + return new TckResult("structural", NAME, TckStatus.FAIL, "BTRACE-LO-001", e.getMessage()); + } + } + + private Path extractApiJar(Path extensionZip) throws IOException { + try (var zf = new ZipFile(extensionZip.toFile())) { + var entry = Collections.list(zf.entries()).stream() + .filter(e -> e.getName().endsWith("-api.jar")).findFirst().orElse(null); + if (entry == null) return null; + var tmp = Files.createTempFile("btrace-tck-api-", ".jar"); + tmp.toFile().deleteOnExit(); + try (var in = zf.getInputStream(entry); + var out = new FileOutputStream(tmp.toFile())) { + in.transferTo(out); + } + return tmp; + } + } +} +``` + +`ApiImplPartitionCheck.java`: +```java +package io.btrace.tck.check.structural; + +import io.btrace.ext.validator.*; +import io.btrace.tck.*; +import java.io.*; +import java.nio.file.*; +import java.util.*; +import java.util.jar.*; +import java.util.zip.*; + +public final class ApiImplPartitionCheck { + private static final String NAME = "ApiImplPartitionCheck"; + + public TckResult run(Path extensionZip) { + try (var zf = new ZipFile(extensionZip.toFile())) { + var entries = Collections.list(zf.entries()); + var apiEntry = entries.stream().filter(e -> e.getName().endsWith("-api.jar")).findFirst().orElse(null); + var implEntry = entries.stream().filter(e -> e.getName().endsWith("-impl.jar")).findFirst().orElse(null); + if (apiEntry == null) + return new TckResult("structural", NAME, TckStatus.FAIL, "BTRACE-LO-002", "No api.jar"); + if (implEntry == null) + return new TckResult("structural", NAME, TckStatus.FAIL, "BTRACE-LO-004", "No impl.jar"); + + Path apiTmp = extract(zf, apiEntry); + Path implTmp = extract(zf, implEntry); + var results = new ApiImplPartitionValidator().validate(apiTmp, implTmp); + var first = results.stream().filter(ValidationResult::isError).findFirst(); + if (first.isPresent()) + return new TckResult("structural", NAME, TckStatus.FAIL, + first.get().getRuleCode(), first.get().getMessage()); + return new TckResult("structural", NAME, TckStatus.PASS, null, null); + } catch (IOException e) { + return new TckResult("structural", NAME, TckStatus.FAIL, "BTRACE-LO-001", e.getMessage()); + } + } + + private Path extract(ZipFile zf, ZipEntry entry) throws IOException { + var tmp = Files.createTempFile("btrace-tck-", ".jar"); + tmp.toFile().deleteOnExit(); + try (var in = zf.getInputStream(entry); var out = new FileOutputStream(tmp.toFile())) { + in.transferTo(out); + } + return tmp; + } +} +``` + +`StructuralSuite.java`: +```java +package io.btrace.tck.suite; + +import io.btrace.tck.*; +import io.btrace.tck.check.structural.*; +import java.util.List; + +public final class StructuralSuite implements Suite { + @Override public String name() { return "structural"; } + + @Override public TckSuiteResult run(TckInput input) { + if (input.extensionZip() == null) { + return new TckSuiteResult("structural", List.of( + new TckResult("structural", "setup", TckStatus.FAIL, + "BTRACE-LO-001", "--extension ZIP path is required"))); + } + var zip = input.extensionZip(); + return new TckSuiteResult("structural", List.of( + new ManifestCheck().run(zip), + new ArtifactLayoutCheck().run(zip), + new ServiceApiCheck().run(zip), + new ApiImplPartitionCheck().run(zip))); + } +} +``` + +- [ ] **Step 3: Run tests — expect PASS** + +```bash +./gradlew :btrace-tck:test +``` + +- [ ] **Step 4: Commit** + +```bash +git add btrace-tck/ +git commit -m "feat(tck): add StructuralSuite with ManifestCheck, ArtifactLayoutCheck, ServiceApiCheck, ApiImplPartitionCheck" +``` + +--- + +## Task 10: Fixture extension ZIPs and fixture-based integration tests + +**Files:** +- Create: fixture build script to generate `good-extension.zip`, `bad-manifest.zip`, `bad-api.zip`, `bad-partition.zip` under `btrace-tck/src/test/resources/fixtures/` +- Create: `btrace-tck/src/test/java/io/btrace/tck/suite/StructuralSuiteFixtureTest.java` + +The simplest approach is to generate the fixture ZIPs programmatically in a JUnit `@BeforeAll` method using ASM — avoids a separate build step and keeps fixtures co-located with the tests. + +- [ ] **Step 1: Create fixture builder helper** + +Create `btrace-tck/src/test/java/io/btrace/tck/suite/FixtureBuilder.java`: +```java +package io.btrace.tck.suite; + +import org.objectweb.asm.*; +import org.objectweb.asm.Opcodes; +import java.io.*; +import java.nio.file.*; +import java.util.*; +import java.util.jar.*; +import java.util.zip.*; + +final class FixtureBuilder { + /** Builds a fully-valid extension ZIP. */ + static Path goodExtension(Path dir) throws Exception { + return buildZip(dir, "good-extension.zip", + validApiJar(dir, "io/example/MyService"), + "good-extension-impl.jar", true); + } + + static Path badManifest(Path dir) throws Exception { + // api.jar with BTrace-Extension-Id missing + var apiJar = dir.resolve("bad-manifest-api.jar"); + try (var jos = new JarOutputStream(new FileOutputStream(apiJar.toFile()))) { + var mf = new Manifest(); + mf.getMainAttributes().put(Attributes.Name.MANIFEST_VERSION, "1.0"); + // deliberately omit BTrace-Extension-Id + mf.getMainAttributes().putValue("BTrace-Extension-Version", "1.0.0"); + mf.getMainAttributes().putValue("BTrace-API-Version", "3.0+"); + mf.getMainAttributes().putValue("BTrace-Extension-Services", "io.example.MyService"); + mf.getMainAttributes().putValue("BTrace-Extension-Impl", "bad-manifest-impl.jar"); + jos.putNextEntry(new JarEntry("META-INF/MANIFEST.MF")); mf.write(jos); jos.closeEntry(); + } + return buildZip(dir, "bad-manifest.zip", apiJar, "bad-manifest-impl.jar", true); + } + + static Path badApi(Path dir) throws Exception { + // api.jar whose service is a class, not an interface → BTRACE-EXT-001 + var apiJar = dir.resolve("bad-api-api.jar"); + var cw = new ClassWriter(0); + cw.visit(Opcodes.V11, Opcodes.ACC_PUBLIC, "io/example/NotAnInterface", null, "java/lang/Object", null); + cw.visitEnd(); + try (var jos = new JarOutputStream(new FileOutputStream(apiJar.toFile()))) { + var mf = new Manifest(); + mf.getMainAttributes().put(Attributes.Name.MANIFEST_VERSION, "1.0"); + mf.getMainAttributes().putValue("BTrace-Extension-Id", "io.example.bad-api"); + mf.getMainAttributes().putValue("BTrace-Extension-Version", "1.0.0"); + mf.getMainAttributes().putValue("BTrace-Extension-Name", "Bad API"); + mf.getMainAttributes().putValue("BTrace-Extension-Description", "bad"); + mf.getMainAttributes().putValue("BTrace-API-Version", "3.0+"); + mf.getMainAttributes().putValue("BTrace-Java-Version", "8+"); + mf.getMainAttributes().putValue("BTrace-Extension-Services", "io.example.NotAnInterface"); + mf.getMainAttributes().putValue("BTrace-Extension-Impl", "bad-api-impl.jar"); + jos.putNextEntry(new JarEntry("META-INF/MANIFEST.MF")); mf.write(jos); jos.closeEntry(); + jos.putNextEntry(new JarEntry("io/example/NotAnInterface.class")); jos.write(cw.toByteArray()); jos.closeEntry(); + } + return buildZip(dir, "bad-api.zip", apiJar, "bad-api-impl.jar", true); + } + + static Path badPartition(Path dir) throws Exception { + // api.jar service method returns an impl-only type → BTRACE-PT-001 + var implClass = "io/example/internal/Impl"; + var svcClass = "io/example/LeakyService"; + var cw = new ClassWriter(0); + cw.visit(Opcodes.V11, Opcodes.ACC_PUBLIC | Opcodes.ACC_INTERFACE | Opcodes.ACC_ABSTRACT, + svcClass, null, "java/lang/Object", null); + cw.visitMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_ABSTRACT, + "get", "()L" + implClass + ";", null, null).visitEnd(); + cw.visitEnd(); + + var implCw = new ClassWriter(0); + implCw.visit(Opcodes.V11, Opcodes.ACC_PUBLIC, implClass, null, "java/lang/Object", null); + implCw.visitEnd(); + + var apiJar = dir.resolve("bad-partition-api.jar"); + var implJar = dir.resolve("bad-partition-impl.jar"); + try (var jos = new JarOutputStream(new FileOutputStream(apiJar.toFile()))) { + var mf = buildFullManifest("io.example.bad-partition", svcClass.replace('/', '.'), "bad-partition-impl.jar"); + jos.putNextEntry(new JarEntry("META-INF/MANIFEST.MF")); mf.write(jos); jos.closeEntry(); + jos.putNextEntry(new JarEntry(svcClass + ".class")); jos.write(cw.toByteArray()); jos.closeEntry(); + } + try (var jos = new JarOutputStream(new FileOutputStream(implJar.toFile()))) { + jos.putNextEntry(new JarEntry(implClass + ".class")); jos.write(implCw.toByteArray()); jos.closeEntry(); + } + return buildZip2(dir, "bad-partition.zip", apiJar, implJar); + } + + private static Path validApiJar(Path dir, String svcInternal) throws Exception { + var cw = new ClassWriter(0); + cw.visit(Opcodes.V11, Opcodes.ACC_PUBLIC | Opcodes.ACC_INTERFACE | Opcodes.ACC_ABSTRACT, + svcInternal, null, "java/lang/Object", null); + var mv = cw.visitMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_ABSTRACT, + "doWork", "()Ljava/lang/String;", null, null); + mv.visitAnnotation("Ljavax/annotation/Nullable;", true).visitEnd(); + mv.visitEnd(); + cw.visitEnd(); + var apiJar = dir.resolve("good-extension-api.jar"); + var svcFqcn = svcInternal.replace('/', '.'); + try (var jos = new JarOutputStream(new FileOutputStream(apiJar.toFile()))) { + var mf = buildFullManifest("io.example.good-extension", svcFqcn, "good-extension-impl.jar"); + jos.putNextEntry(new JarEntry("META-INF/MANIFEST.MF")); mf.write(jos); jos.closeEntry(); + jos.putNextEntry(new JarEntry(svcInternal + ".class")); jos.write(cw.toByteArray()); jos.closeEntry(); + } + return apiJar; + } + + private static Manifest buildFullManifest(String id, String services, String impl) { + var mf = new Manifest(); + var a = mf.getMainAttributes(); + a.put(Attributes.Name.MANIFEST_VERSION, "1.0"); + a.putValue("BTrace-Extension-Id", id); + a.putValue("BTrace-Extension-Version", "1.0.0"); + a.putValue("BTrace-Extension-Name", id); + a.putValue("BTrace-Extension-Description", "Test extension"); + a.putValue("BTrace-API-Version", "3.0+"); + a.putValue("BTrace-Java-Version", "8+"); + a.putValue("BTrace-Extension-Services", services); + a.putValue("BTrace-Extension-Impl", impl); + return mf; + } + + private static Path buildZip(Path dir, String name, Path apiJar, String implName, boolean createImpl) throws Exception { + var implJar = dir.resolve(implName); + if (createImpl) implJar.toFile().createNewFile(); + return buildZip2(dir, name, apiJar, implJar); + } + + private static Path buildZip2(Path dir, String name, Path apiJar, Path implJar) throws Exception { + var zip = dir.resolve(name); + try (var zos = new ZipOutputStream(new FileOutputStream(zip.toFile()))) { + for (Path p : List.of(apiJar, implJar)) { + zos.putNextEntry(new ZipEntry(p.getFileName().toString())); + Files.copy(p, zos); + zos.closeEntry(); + } + } + return zip; + } +} +``` + +- [ ] **Step 2: Write fixture integration tests** + +Create `btrace-tck/src/test/java/io/btrace/tck/suite/StructuralSuiteFixtureTest.java`: +```java +package io.btrace.tck.suite; + +import io.btrace.tck.*; +import org.junit.jupiter.api.*; +import org.junit.jupiter.api.io.TempDir; +import java.nio.file.Path; +import static org.junit.jupiter.api.Assertions.*; + +class StructuralSuiteFixtureTest { + @TempDir static Path tmp; + static Path goodZip, badManifestZip, badApiZip, badPartitionZip; + + @BeforeAll static void buildFixtures() throws Exception { + goodZip = FixtureBuilder.goodExtension(tmp); + badManifestZip = FixtureBuilder.badManifest(tmp); + badApiZip = FixtureBuilder.badApi(tmp); + badPartitionZip = FixtureBuilder.badPartition(tmp); + } + + private TckSuiteResult run(Path zip) { + return new StructuralSuite().run( + TckInput.builder().extensionZip(zip).reportDir(tmp.resolve("report")).build()); + } + + @Test void goodExtensionPasses() { + var result = run(goodZip); + assertFalse(result.hasFailed(), () -> "Expected PASS but got: " + result.checks()); + } + + @Test void badManifestFails() { + var result = run(badManifestZip); + assertTrue(result.hasFailed()); + assertTrue(result.checks().stream().anyMatch(r -> + r.isFail() && r.ruleCode() != null && r.ruleCode().startsWith("BTRACE-MF-"))); + } + + @Test void badApiFails() { + var result = run(badApiZip); + assertTrue(result.hasFailed()); + assertTrue(result.checks().stream().anyMatch(r -> + r.isFail() && "BTRACE-EXT-001".equals(r.ruleCode()))); + } + + @Test void badPartitionFails() { + var result = run(badPartitionZip); + assertTrue(result.hasFailed()); + assertTrue(result.checks().stream().anyMatch(r -> + r.isFail() && "BTRACE-PT-001".equals(r.ruleCode()))); + } +} +``` + +The `FixtureBuilder` needs ASM on the test classpath. Add to `btrace-tck/build.gradle`: +```groovy +testImplementation 'org.ow2.asm:asm:9.9.1' +testImplementation 'org.ow2.asm:asm-tree:9.9.1' +``` + +- [ ] **Step 3: Run tests — expect PASS** + +```bash +./gradlew :btrace-tck:test +``` +Expected: `BUILD SUCCESSFUL`, all fixture tests pass. + +- [ ] **Step 4: Commit** + +```bash +git add btrace-tck/ +git commit -m "feat(tck): add FixtureBuilder and StructuralSuiteFixtureTest" +``` + +--- + +## Task 11: TckMain CLI and end-to-end test + +**Files:** +- Create: `btrace-tck/src/main/java/io/btrace/tck/cli/TckMain.java` +- Create: `btrace-tck/src/test/java/io/btrace/tck/cli/TckMainTest.java` + +- [ ] **Step 1: Write failing test** + +Create `btrace-tck/src/test/java/io/btrace/tck/cli/TckMainTest.java`: +```java +package io.btrace.tck.cli; + +import io.btrace.tck.suite.FixtureBuilder; +import org.junit.jupiter.api.*; +import org.junit.jupiter.api.io.TempDir; +import java.nio.file.*; +import static org.junit.jupiter.api.Assertions.*; + +class TckMainTest { + @TempDir static Path tmp; + + @Test void exitCode0ForGoodExtension() throws Exception { + var zip = FixtureBuilder.goodExtension(tmp); + var report = tmp.resolve("report"); + int code = TckMain.run(new String[]{ + "--extension", zip.toString(), + "--report-dir", report.toString(), + "--suites", "structural" + }); + assertEquals(0, code, "Expected exit 0 for valid extension"); + assertTrue(Files.exists(report.resolve("tck-results.xml"))); + assertTrue(Files.exists(report.resolve("tck-report.html"))); + } + + @Test void exitCode1ForBadManifest() throws Exception { + var zip = FixtureBuilder.badManifest(tmp.resolve("bm")); + Files.createDirectories(tmp.resolve("bm")); + var report = tmp.resolve("report-bm"); + int code = TckMain.run(new String[]{ + "--extension", zip.toString(), + "--report-dir", report.toString(), + "--suites", "structural" + }); + assertEquals(1, code, "Expected exit 1 for invalid extension"); + } +} +``` + +- [ ] **Step 2: Implement TckMain** + +Create `btrace-tck/src/main/java/io/btrace/tck/cli/TckMain.java`: +```java +package io.btrace.tck.cli; + +import io.btrace.tck.*; +import io.btrace.tck.report.*; +import io.btrace.tck.suite.StructuralSuite; +import picocli.CommandLine; +import picocli.CommandLine.Command; +import picocli.CommandLine.Option; +import java.io.PrintWriter; +import java.nio.file.Path; +import java.util.*; +import java.util.concurrent.Callable; + +@Command(name = "btrace-tck", mixinStandardHelpOptions = true, + description = "BTrace Extension Technology Compatibility Kit") +public final class TckMain implements Callable { + + @Option(names = {"--extension", "-e"}, required = true, description = "Path to extension ZIP") + Path extensionZip; + + @Option(names = {"--btrace-home"}, description = "Path to BTrace installation") + Path btraceHome; + + @Option(names = {"--report-dir", "-r"}, defaultValue = "tck-report", description = "Output directory") + Path reportDir; + + @Option(names = {"--suites", "-s"}, split = ",", defaultValue = "structural", + description = "Suites to run: structural,lifecycle,behavioral,perf") + List suites; + + public static void main(String[] args) { + System.exit(run(args)); + } + + public static int run(String[] args) { + return new CommandLine(new TckMain()).execute(args); + } + + @Override + public Integer call() throws Exception { + var input = TckInput.builder() + .extensionZip(extensionZip) + .btraceHome(btraceHome) + .reportDir(reportDir) + .build(); + + var activeSuites = new ArrayList(); + for (String s : suites) { + if ("structural".equals(s)) activeSuites.add(new StructuralSuite()); + // lifecycle, behavioral, perf suites added in Plans 2 and 3 + } + + var engine = new TckEngine(activeSuites); + var results = engine.run(input); + + reportDir.toFile().mkdirs(); + new StdoutReporter(new PrintWriter(System.out, true)).report(results); + new JUnitXmlReporter(reportDir.resolve("tck-results.xml")).report(results); + new HtmlReporter(reportDir.resolve("tck-report.html")).report(results); + + boolean failed = results.stream().anyMatch(TckSuiteResult::hasFailed); + return failed ? 1 : 0; + } +} +``` + +- [ ] **Step 3: Run tests — expect PASS** + +```bash +./gradlew :btrace-tck:test +``` + +- [ ] **Step 4: Verify standalone JAR is runnable** + +```bash +./gradlew :btrace-tck:jar +``` +Expected: `btrace-tck/build/libs/btrace-tck.jar` created. (No need to run it now — the TckMainTest covers the behavior.) + +- [ ] **Step 5: Spotless** + +```bash +./gradlew spotlessApply +``` + +- [ ] **Step 6: Commit** + +```bash +git add btrace-tck/ btrace-ext-validator/ +git commit -m "feat(tck): add TckMain CLI and end-to-end structural test" +``` + +--- + +## Task 12: Run full build and format check + +- [ ] **Step 1: Run all new module tests** + +```bash +./gradlew :btrace-ext-validator:test :btrace-tck:test +``` +Expected: all tests pass. + +- [ ] **Step 2: Run spotlessCheck** + +```bash +./gradlew spotlessCheck +``` +If any formatting violations: +```bash +./gradlew spotlessApply +git add -u +git commit -m "style: apply google-java-format to tck modules" +``` + +- [ ] **Step 3: Run btrace-gradle-plugin tests to verify refactor didn't break anything** + +```bash +./gradlew :btrace-gradle-plugin:test +``` + +- [ ] **Step 4: Smoke-test validateServiceApis on existing extensions** + +```bash +./gradlew :btrace-extensions:btrace-contracts:validateServiceApis \ + :btrace-extensions:btrace-metrics:validateServiceApis +``` +Expected: no new failures. + +- [ ] **Step 5: Final commit if any cleanup needed** + +```bash +git add -u +git commit -m "chore(tck): cleanup after Plan 1 integration" +``` + +--- + +## Self-Review Notes + +- **Spec coverage:** All structural checks from the spec are implemented (ManifestCheck, ArtifactLayoutCheck, ServiceApiCheck, ApiImplPartitionCheck). RegistryCheck (optional) is omitted from Plan 1 — it's low-priority and standalone schema validation. Add it as a follow-up task if needed. +- **BTRACE-EXT-022 nuance:** The `ServiceApiValidator.checkNullability` method applies BTRACE-EXT-022 to all Object return types, not just interface return types. The distinction requires loading the class (to call `isInterface()`) or keeping a set of known interface names from the JAR. The current implementation conservatively applies it to all Object returns — this is safe (stricter than required) and matches the plugin's behavior. Revisit if it causes false positives on extension authors returning concrete non-interface types. +- **Plan 2** covers: LifecycleSuite (LoadCheck, InitCheck, InjectionCheck, CloseCheck) and BehavioralSuite — requires spawning a child JVM with the BTrace agent. +- **Plan 3** covers: PerformanceSuite (JMH) and btrace-tck-gradle-plugin. diff --git a/internal/superpowers/plans/2026-05-17-extension-tck-plan-2-lifecycle-behavioral.md b/internal/superpowers/plans/2026-05-17-extension-tck-plan-2-lifecycle-behavioral.md new file mode 100644 index 000000000..6be9098b0 --- /dev/null +++ b/internal/superpowers/plans/2026-05-17-extension-tck-plan-2-lifecycle-behavioral.md @@ -0,0 +1,1214 @@ +# Extension TCK — Part 2: Lifecycle & Behavioral Suites Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add `LifecycleSuite` and `BehavioralSuite` to `btrace-tck`, verifying that an extension loads cleanly, initialises, makes services injectable, falls back gracefully to shims when denied, and remains data-race-free under concurrent access — all without spawning a child JVM. + +**Architecture:** Both suites use `ExtensionLoaderImpl` and `ExtensionBridgeImpl` from `btrace-core` directly in the TCK's own JVM. `LifecycleSuite` checks the load→init→inject→close sequence. `BehavioralSuite` reuses the same loader session to test shim fallback (via `PermissionPolicy`), null safety of NoOp shims, concurrency, and required-service failure mode. `TckMain` is updated to wire these suites in after structural passes. + +**Tech Stack:** btrace-core (`ExtensionLoaderImpl`, `ExtensionBridgeImpl`, `PermissionPolicy`, `FileSystemExtensionRepository`), Java 11 concurrency primitives, JUnit 5 for self-tests. + +**Prerequisite:** Plan 1 complete (`btrace-ext-validator` built and all structural checks passing). + +--- + +## File Map + +### Modified: `btrace-tck/` +| File | Change | +|------|--------| +| `build.gradle` | Add `implementation project(':btrace-core')` | +| `src/main/java/io/btrace/tck/suite/LifecycleSuite.java` | New suite | +| `src/main/java/io/btrace/tck/check/lifecycle/LoadCheck.java` | Extension loads without error | +| `src/main/java/io/btrace/tck/check/lifecycle/InitCheck.java` | `isLoaded()` true after `load()` | +| `src/main/java/io/btrace/tck/check/lifecycle/InjectionCheck.java` | `bridge.getExtensionClass()` non-null | +| `src/main/java/io/btrace/tck/check/lifecycle/CloseCheck.java` | `loader.close()` without exception | +| `src/main/java/io/btrace/tck/suite/BehavioralSuite.java` | New suite | +| `src/main/java/io/btrace/tck/check/behavioral/NullSafetyCheck.java` | NoOp shim methods don't throw | +| `src/main/java/io/btrace/tck/check/behavioral/ShimFallbackCheck.java` | Denied ext → shim substituted | +| `src/main/java/io/btrace/tck/check/behavioral/ConcurrencyCheck.java` | 8-thread × 1000 calls, no races | +| `src/main/java/io/btrace/tck/check/behavioral/RequiredServiceCheck.java` | Required service fails fast if denied | +| `src/main/java/io/btrace/tck/cli/TckMain.java` | Wire in lifecycle and behavioral suites | +| `src/main/java/io/btrace/tck/ExtensionSession.java` | Shared lifecycle state for both suites | +| `src/test/java/io/btrace/tck/suite/LifecycleSuiteTest.java` | Unit tests | +| `src/test/java/io/btrace/tck/suite/BehavioralSuiteTest.java` | Unit tests | + +--- + +## Task 1: Add btrace-core dependency + +**Files:** +- Modify: `btrace-tck/build.gradle` + +- [ ] **Step 1: Add dependency** + +In `btrace-tck/build.gradle`, add to `dependencies`: +```groovy +implementation project(':btrace-core') +``` + +- [ ] **Step 2: Verify it resolves** + +```bash +./gradlew :btrace-tck:dependencies --configuration runtimeClasspath 2>&1 | grep btrace-core +``` +Expected: `+--- project :btrace-core` in output. + +- [ ] **Step 3: Run existing tests to verify no regression** + +```bash +./gradlew :btrace-tck:test +``` +Expected: `BUILD SUCCESSFUL`. + +- [ ] **Step 4: Commit** + +```bash +git add btrace-tck/build.gradle +git commit -m "feat(tck): add btrace-core dependency to btrace-tck" +``` + +--- + +## Task 2: ExtensionSession — shared lifecycle state + +Both suites need to open an extension, use it, and close it. `ExtensionSession` encapsulates this lifecycle so each suite doesn't duplicate the setup/teardown logic. + +**Files:** +- Create: `btrace-tck/src/main/java/io/btrace/tck/ExtensionSession.java` +- Create: `btrace-tck/src/test/java/io/btrace/tck/ExtensionSessionTest.java` + +- [ ] **Step 1: Write failing test** + +Create `btrace-tck/src/test/java/io/btrace/tck/ExtensionSessionTest.java`: +```java +package io.btrace.tck; + +import io.btrace.tck.suite.FixtureBuilder; +import org.junit.jupiter.api.*; +import org.junit.jupiter.api.io.TempDir; +import java.nio.file.Path; +import static org.junit.jupiter.api.Assertions.*; + +class ExtensionSessionTest { + @TempDir static Path tmp; + static Path goodZip; + + @BeforeAll static void buildFixtures() throws Exception { + goodZip = FixtureBuilder.goodExtension(tmp); + } + + @Test void openAndCloseWithoutErrors() throws Exception { + try (var session = ExtensionSession.open(goodZip, tmp.resolve("ext-dir"))) { + assertNotNull(session.descriptor()); + assertTrue(session.descriptor().isLoaded()); + } + } + + @Test void openFailsForNonExistentZip() { + assertThrows(ExtensionSession.LoadException.class, + () -> ExtensionSession.open(tmp.resolve("nonexistent.zip"), tmp.resolve("x"))); + } +} +``` + +- [ ] **Step 2: Run test — expect compilation failure** + +```bash +./gradlew :btrace-tck:test 2>&1 | grep "error:" +``` + +- [ ] **Step 3: Implement ExtensionSession** + +Create `btrace-tck/src/main/java/io/btrace/tck/ExtensionSession.java`: +```java +package io.btrace.tck; + +import io.btrace.extension.ExtensionLoader; +import io.btrace.extension.ExtensionDescriptorDTO; +import io.btrace.extension.impl.ExtensionBridgeImpl; +import io.btrace.extension.impl.ExtensionConfig; +import io.btrace.extension.impl.ExtensionLoaderImpl; +import io.btrace.extension.impl.FileSystemExtensionRepository; +import java.io.*; +import java.nio.file.*; +import java.util.*; +import java.util.zip.ZipFile; + +/** + * Owns the full lifecycle of an extension under test: + * extract ZIP → create loader → discover → load → (use) → close. + * + * Implements AutoCloseable so callers can use try-with-resources. + */ +public final class ExtensionSession implements AutoCloseable { + + private final ExtensionLoaderImpl loader; + private final ExtensionBridgeImpl bridge; + private final ExtensionDescriptorDTO descriptor; + + private ExtensionSession(ExtensionLoaderImpl loader, + ExtensionBridgeImpl bridge, + ExtensionDescriptorDTO descriptor) { + this.loader = loader; + this.bridge = bridge; + this.descriptor = descriptor; + } + + /** + * Open an extension session from an extension ZIP. + * + * @param extensionZip path to *-extension.zip + * @param workDir temp directory to extract the ZIP contents into + */ + public static ExtensionSession open(Path extensionZip, Path workDir) throws LoadException { + if (!Files.exists(extensionZip)) { + throw new LoadException("Extension ZIP not found: " + extensionZip); + } + try { + Files.createDirectories(workDir); + extractZip(extensionZip, workDir); + + var repo = new FileSystemExtensionRepository(workDir, 200 /* Priority.USER */); + var loader = new ExtensionLoaderImpl( + List.of(repo), + ExtensionSession.class.getClassLoader(), + ExtensionConfig.createDefault(), + null, // no Instrumentation needed for unit-level checks + "tck-test"); + + var extensions = loader.discoverExtensions(); + if (extensions.isEmpty()) { + throw new LoadException("No extensions discovered in: " + extensionZip); + } + var desc = extensions.get(0); + loader.load(desc); + + var bridge = new ExtensionBridgeImpl(loader); + return new ExtensionSession(loader, bridge, desc); + } catch (LoadException e) { + throw e; + } catch (Exception e) { + throw new LoadException("Failed to open extension session: " + e.getMessage(), e); + } + } + + public ExtensionDescriptorDTO descriptor() { return descriptor; } + public ExtensionLoaderImpl loader() { return loader; } + public ExtensionBridgeImpl bridge() { return bridge; } + + /** + * Resolve the implementation class for a declared service. + * Returns null if the service is not available (e.g., denied by policy). + */ + public Class resolveService(String serviceClassName) { + try { + return bridge.getExtensionClass(serviceClassName); + } catch (Exception e) { + return null; + } + } + + @Override + public void close() { + try { loader.close(); } catch (Exception ignored) {} + } + + private static void extractZip(Path zip, Path targetDir) throws IOException { + try (var zf = new ZipFile(zip.toFile())) { + var entries = Collections.list(zf.entries()); + for (var entry : entries) { + var dest = targetDir.resolve(entry.getName()); + if (entry.isDirectory()) { + Files.createDirectories(dest); + } else { + Files.createDirectories(dest.getParent()); + try (var in = zf.getInputStream(entry); + var out = new FileOutputStream(dest.toFile())) { + in.transferTo(out); + } + } + } + } + } + + public static final class LoadException extends Exception { + public LoadException(String message) { super(message); } + public LoadException(String message, Throwable cause) { super(message, cause); } + } +} +``` + +- [ ] **Step 4: Run tests — expect PASS** + +```bash +./gradlew :btrace-tck:test +``` + +- [ ] **Step 5: Commit** + +```bash +git add btrace-tck/ +git commit -m "feat(tck): add ExtensionSession for lifecycle state management" +``` + +--- + +## Task 3: LifecycleSuite + +**Files:** +- Create: `btrace-tck/src/main/java/io/btrace/tck/check/lifecycle/LoadCheck.java` +- Create: `btrace-tck/src/main/java/io/btrace/tck/check/lifecycle/InitCheck.java` +- Create: `btrace-tck/src/main/java/io/btrace/tck/check/lifecycle/InjectionCheck.java` +- Create: `btrace-tck/src/main/java/io/btrace/tck/check/lifecycle/CloseCheck.java` +- Create: `btrace-tck/src/main/java/io/btrace/tck/suite/LifecycleSuite.java` +- Create: `btrace-tck/src/test/java/io/btrace/tck/suite/LifecycleSuiteTest.java` + +- [ ] **Step 1: Write failing test** + +Create `btrace-tck/src/test/java/io/btrace/tck/suite/LifecycleSuiteTest.java`: +```java +package io.btrace.tck.suite; + +import io.btrace.tck.*; +import org.junit.jupiter.api.*; +import org.junit.jupiter.api.io.TempDir; +import java.nio.file.Path; +import static org.junit.jupiter.api.Assertions.*; + +class LifecycleSuiteTest { + @TempDir static Path tmp; + static Path goodZip; + + @BeforeAll static void buildFixtures() throws Exception { + goodZip = FixtureBuilder.goodExtension(tmp); + } + + @Test void suiteIsNamedLifecycle() { + assertEquals("lifecycle", new LifecycleSuite().name()); + } + + @Test void goodExtensionPassesAllLifecycleChecks() { + var input = TckInput.builder() + .extensionZip(goodZip) + .reportDir(tmp.resolve("report")) + .build(); + var result = new LifecycleSuite().run(input); + assertFalse(result.hasFailed(), () -> "Lifecycle checks failed: " + result.checks()); + } + + @Test void missingZipProducesFailResult() { + var input = TckInput.builder() + .extensionZip(tmp.resolve("nonexistent.zip")) + .reportDir(tmp.resolve("report")) + .build(); + var result = new LifecycleSuite().run(input); + assertTrue(result.hasFailed()); + } +} +``` + +- [ ] **Step 2: Implement the four checks** + +`LoadCheck.java`: +```java +package io.btrace.tck.check.lifecycle; + +import io.btrace.tck.*; +import java.nio.file.Path; + +public final class LoadCheck { + private static final String NAME = "LoadCheck"; + + public TckResult run(Path extensionZip, Path workDir) { + try (var session = ExtensionSession.open(extensionZip, workDir)) { + var desc = session.descriptor(); + if (!desc.isLoaded()) { + return new TckResult("lifecycle", NAME, TckStatus.FAIL, + "BTRACE-LC-001", "Extension loaded=false after load()"); + } + return new TckResult("lifecycle", NAME, TckStatus.PASS, null, null); + } catch (ExtensionSession.LoadException e) { + return new TckResult("lifecycle", NAME, TckStatus.FAIL, + "BTRACE-LC-001", "Extension failed to load: " + e.getMessage()); + } + } +} +``` + +`InitCheck.java`: +```java +package io.btrace.tck.check.lifecycle; + +import io.btrace.tck.*; +import java.nio.file.Path; + +public final class InitCheck { + private static final String NAME = "InitCheck"; + + public TckResult run(Path extensionZip, Path workDir) { + try (var session = ExtensionSession.open(extensionZip, workDir)) { + var desc = session.descriptor(); + // isLoaded() becomes true only after initialize() completes inside load() + if (!desc.isLoaded()) { + return new TckResult("lifecycle", NAME, TckStatus.FAIL, + "BTRACE-LC-002", "Extension initialize() did not complete (isLoaded=false)"); + } + return new TckResult("lifecycle", NAME, TckStatus.PASS, null, null); + } catch (ExtensionSession.LoadException e) { + return new TckResult("lifecycle", NAME, TckStatus.FAIL, + "BTRACE-LC-002", "Load failed before initialize(): " + e.getMessage()); + } + } +} +``` + +`InjectionCheck.java`: +```java +package io.btrace.tck.check.lifecycle; + +import io.btrace.tck.*; +import java.nio.file.Path; +import java.util.List; + +public final class InjectionCheck { + private static final String NAME = "InjectionCheck"; + + public TckResult run(Path extensionZip, Path workDir) { + try (var session = ExtensionSession.open(extensionZip, workDir)) { + List services = session.descriptor().getServices(); + if (services.isEmpty()) { + return new TckResult("lifecycle", NAME, TckStatus.PASS, + null, "No services declared; skipping injection check"); + } + for (String svc : services) { + Class resolved = session.resolveService(svc); + if (resolved == null) { + return new TckResult("lifecycle", NAME, TckStatus.FAIL, + "BTRACE-LC-003", "Service '" + svc + "' could not be injected (null)"); + } + // Verify it's not just the raw interface (should be an impl or shim) + if (resolved.isInterface()) { + return new TckResult("lifecycle", NAME, TckStatus.FAIL, + "BTRACE-LC-003", + "Service '" + svc + "' resolved to interface only — impl or shim expected"); + } + } + return new TckResult("lifecycle", NAME, TckStatus.PASS, null, null); + } catch (ExtensionSession.LoadException e) { + return new TckResult("lifecycle", NAME, TckStatus.FAIL, + "BTRACE-LC-003", "Load failed before injection: " + e.getMessage()); + } + } +} +``` + +`CloseCheck.java`: +```java +package io.btrace.tck.check.lifecycle; + +import io.btrace.tck.*; +import java.nio.file.Path; + +public final class CloseCheck { + private static final String NAME = "CloseCheck"; + + public TckResult run(Path extensionZip, Path workDir) { + try { + var session = ExtensionSession.open(extensionZip, workDir); + // close() must not throw + session.close(); + return new TckResult("lifecycle", NAME, TckStatus.PASS, null, null); + } catch (ExtensionSession.LoadException e) { + return new TckResult("lifecycle", NAME, TckStatus.FAIL, + "BTRACE-LC-004", "Load failed before close: " + e.getMessage()); + } catch (Exception e) { + return new TckResult("lifecycle", NAME, TckStatus.FAIL, + "BTRACE-LC-004", "close() threw: " + e.getMessage()); + } + } +} +``` + +`LifecycleSuite.java`: +```java +package io.btrace.tck.suite; + +import io.btrace.tck.*; +import io.btrace.tck.check.lifecycle.*; +import java.io.IOException; +import java.nio.file.*; +import java.util.List; + +public final class LifecycleSuite implements Suite { + @Override public String name() { return "lifecycle"; } + + @Override public TckSuiteResult run(TckInput input) { + if (input.extensionZip() == null) { + return fail("--extension ZIP path is required"); + } + Path workDir; + try { + workDir = Files.createTempDirectory("btrace-tck-lifecycle-"); + } catch (IOException e) { + return fail("Cannot create work dir: " + e.getMessage()); + } + var zip = input.extensionZip(); + return new TckSuiteResult("lifecycle", List.of( + new LoadCheck().run(zip, workDir.resolve("load")), + new InitCheck().run(zip, workDir.resolve("init")), + new InjectionCheck().run(zip, workDir.resolve("inject")), + new CloseCheck().run(zip, workDir.resolve("close")))); + } + + private TckSuiteResult fail(String msg) { + return new TckSuiteResult("lifecycle", List.of( + new TckResult("lifecycle", "setup", TckStatus.FAIL, "BTRACE-LC-001", msg))); + } +} +``` + +- [ ] **Step 3: Run tests — expect PASS** + +```bash +./gradlew :btrace-tck:test --tests "*LifecycleSuite*" +``` +Expected: all lifecycle tests pass. + +- [ ] **Step 4: Commit** + +```bash +git add btrace-tck/ +git commit -m "feat(tck): add LifecycleSuite (LoadCheck, InitCheck, InjectionCheck, CloseCheck)" +``` + +--- + +## Task 4: BehavioralSuite — NullSafetyCheck + +The NoOp shim for each service is stored in the api.jar under `META-INF/btrace/shims.index`. `NullSafetyCheck` loads the NoOp shim class and calls every method via reflection with null/zero/false arguments, verifying no exception escapes to the caller. + +**Files:** +- Create: `btrace-tck/src/main/java/io/btrace/tck/check/behavioral/NullSafetyCheck.java` +- Create: `btrace-tck/src/test/java/io/btrace/tck/check/behavioral/NullSafetyCheckTest.java` + +- [ ] **Step 1: Write failing test** + +Create `btrace-tck/src/test/java/io/btrace/tck/check/behavioral/NullSafetyCheckTest.java`: +```java +package io.btrace.tck.check.behavioral; + +import io.btrace.tck.*; +import io.btrace.tck.suite.FixtureBuilder; +import org.junit.jupiter.api.*; +import org.junit.jupiter.api.io.TempDir; +import java.nio.file.Path; +import static org.junit.jupiter.api.Assertions.*; + +class NullSafetyCheckTest { + @TempDir static Path tmp; + static Path goodZip; + + @BeforeAll static void buildFixtures() throws Exception { + goodZip = FixtureBuilder.goodExtension(tmp); + } + + @Test void goodExtensionShimPassesNullSafety() throws Exception { + var result = new NullSafetyCheck().run(goodZip, tmp.resolve("ns-work")); + // The good-extension fixture's NoOp shim should not throw on any method + assertNotEquals(TckStatus.FAIL, result.status(), + () -> "NullSafetyCheck failed: " + result.message()); + } +} +``` + +- [ ] **Step 2: Implement NullSafetyCheck** + +Create `btrace-tck/src/main/java/io/btrace/tck/check/behavioral/NullSafetyCheck.java`: +```java +package io.btrace.tck.check.behavioral; + +import io.btrace.tck.*; +import java.io.*; +import java.lang.reflect.*; +import java.net.URLClassLoader; +import java.nio.file.*; +import java.util.*; +import java.util.jar.*; +import java.util.zip.*; + +public final class NullSafetyCheck { + private static final String NAME = "NullSafetyCheck"; + + public TckResult run(Path extensionZip, Path workDir) { + try { + Path apiJar = extractApiJar(extensionZip, workDir); + if (apiJar == null) return pass(); // no api.jar — ArtifactLayoutCheck catches this + + var shimsIndex = readShimsIndex(apiJar); + if (shimsIndex.isEmpty()) { + return new TckResult("behavioral", NAME, TckStatus.PASS, + null, "No shims found in api.jar (shim generation may be optional)"); + } + + try (var cl = new URLClassLoader(new java.net.URL[]{apiJar.toUri().toURL()}, + NullSafetyCheck.class.getClassLoader())) { + for (var entry : shimsIndex.entrySet()) { + String shimClass = entry.getValue().get("noop"); + if (shimClass == null) continue; + Class shim = cl.loadClass(shimClass); + Object instance = shim.getDeclaredConstructor().newInstance(); + for (Method m : shim.getDeclaredMethods()) { + if (!Modifier.isPublic(m.getModifiers())) continue; + Object[] args = nullArgs(m.getParameterTypes()); + try { + m.invoke(instance, args); + } catch (InvocationTargetException ite) { + Throwable cause = ite.getCause(); + if (!(cause instanceof UnsupportedOperationException)) { + return new TckResult("behavioral", NAME, TckStatus.FAIL, + "BTRACE-BH-001", + "NoOp shim method '" + m.getName() + "' threw: " + cause); + } + } + } + } + } + return pass(); + } catch (Exception e) { + return new TckResult("behavioral", NAME, TckStatus.FAIL, + "BTRACE-BH-001", "NullSafetyCheck error: " + e.getMessage()); + } + } + + /** + * Reads META-INF/btrace/shims.index from the api.jar. + * Format: one line per service — "serviceClass=noop:NoOpClass,throw:ThrowClass" + */ + private Map> readShimsIndex(Path apiJar) throws Exception { + var result = new LinkedHashMap>(); + try (var jf = new JarFile(apiJar.toFile())) { + var entry = jf.getEntry("META-INF/btrace/shims.index"); + if (entry == null) return result; + try (var reader = new BufferedReader(new InputStreamReader(jf.getInputStream(entry)))) { + String line; + while ((line = reader.readLine()) != null) { + line = line.trim(); + if (line.isEmpty() || line.startsWith("#")) continue; + int eq = line.indexOf('='); + if (eq < 0) continue; + String svc = line.substring(0, eq).trim(); + String pairs = line.substring(eq + 1).trim(); + var map = new HashMap(); + for (String pair : pairs.split(",")) { + int colon = pair.indexOf(':'); + if (colon >= 0) map.put(pair.substring(0, colon).trim(), + pair.substring(colon + 1).trim()); + } + result.put(svc, map); + } + } + } + return result; + } + + private Path extractApiJar(Path extensionZip, Path workDir) throws IOException { + Files.createDirectories(workDir); + try (var zf = new ZipFile(extensionZip.toFile())) { + var entry = Collections.list(zf.entries()).stream() + .filter(e -> e.getName().endsWith("-api.jar")).findFirst().orElse(null); + if (entry == null) return null; + var dest = workDir.resolve(entry.getName()); + try (var in = zf.getInputStream(entry); var out = new FileOutputStream(dest.toFile())) { + in.transferTo(out); + } + return dest; + } + } + + private Object[] nullArgs(Class[] types) { + var args = new Object[types.length]; + for (int i = 0; i < types.length; i++) { + if (types[i] == boolean.class) args[i] = false; + else if (types[i] == int.class + || types[i] == long.class + || types[i] == short.class + || types[i] == byte.class) args[i] = 0; + else if (types[i] == double.class + || types[i] == float.class) args[i] = 0.0; + else if (types[i] == char.class) args[i] = '\0'; + else args[i] = null; + } + return args; + } + + private TckResult pass() { + return new TckResult("behavioral", NAME, TckStatus.PASS, null, null); + } +} +``` + +> **Note on shims.index format:** The exact format written by `generateShimIndex` task in the Gradle plugin must match what `readShimsIndex` parses. Read `BTraceExtensionPlugin.groovy`'s `generateShimIndex` task before finalising this implementation and adjust the parser if the format differs. + +- [ ] **Step 3: Run tests — expect PASS** + +```bash +./gradlew :btrace-tck:test --tests "*NullSafety*" +``` + +- [ ] **Step 4: Commit** + +```bash +git add btrace-tck/ +git commit -m "feat(tck): add NullSafetyCheck for NoOp shim behavioral validation" +``` + +--- + +## Task 5: BehavioralSuite — ShimFallbackCheck and RequiredServiceCheck + +**Files:** +- Create: `btrace-tck/src/main/java/io/btrace/tck/check/behavioral/ShimFallbackCheck.java` +- Create: `btrace-tck/src/main/java/io/btrace/tck/check/behavioral/RequiredServiceCheck.java` +- Create: `btrace-tck/src/test/java/io/btrace/tck/check/behavioral/ShimFallbackCheckTest.java` + +- [ ] **Step 1: Write failing test** + +Create `btrace-tck/src/test/java/io/btrace/tck/check/behavioral/ShimFallbackCheckTest.java`: +```java +package io.btrace.tck.check.behavioral; + +import io.btrace.tck.*; +import io.btrace.tck.suite.FixtureBuilder; +import org.junit.jupiter.api.*; +import org.junit.jupiter.api.io.TempDir; +import java.nio.file.Path; +import static org.junit.jupiter.api.Assertions.*; + +class ShimFallbackCheckTest { + @TempDir static Path tmp; + static Path goodZip; + + @BeforeAll static void buildFixtures() throws Exception { + goodZip = FixtureBuilder.goodExtension(tmp); + } + + @Test void shimFallbackPassesOnGoodExtension() { + var result = new ShimFallbackCheck() + .run(goodZip, tmp.resolve("sf-work"), goodZip /* extensionId from desc */); + assertNotEquals(TckStatus.FAIL, result.status(), () -> result.message()); + } +} +``` + +- [ ] **Step 2: Implement ShimFallbackCheck** + +Create `btrace-tck/src/main/java/io/btrace/tck/check/behavioral/ShimFallbackCheck.java`: +```java +package io.btrace.tck.check.behavioral; + +import io.btrace.extension.PermissionPolicy; +import io.btrace.tck.*; +import java.nio.file.Path; +import java.util.List; + +public final class ShimFallbackCheck { + private static final String NAME = "ShimFallbackCheck"; + + /** + * Deny the extension via policy, then verify that: + * 1. The bridge still resolves a non-null class (the shim/interface). + * 2. The resolved class is NOT the implementation (i.e., shim was substituted). + */ + public TckResult run(Path extensionZip, Path workDir, Path originalZip) { + // Temporarily deny the extension by ID + String extId; + try (var probe = ExtensionSession.open(extensionZip, workDir.resolve("probe"))) { + extId = probe.descriptor().getId(); + } catch (ExtensionSession.LoadException e) { + return new TckResult("behavioral", NAME, TckStatus.FAIL, + "BTRACE-BH-002", "Cannot open extension to read its ID: " + e.getMessage()); + } + + PermissionPolicy.get().setDenyExtensionsCsv(extId); + try (var session = ExtensionSession.open(extensionZip, workDir.resolve("denied"))) { + List services = session.descriptor().getServices(); + for (String svc : services) { + Class resolved = session.resolveService(svc); + if (resolved == null) { + return new TckResult("behavioral", NAME, TckStatus.FAIL, + "BTRACE-BH-002", + "Service '" + svc + "' resolved to null when denied (expected shim or interface)"); + } + // When denied, bridge returns interface or NoOp shim — either is acceptable. + // What's NOT acceptable: it should not be the concrete impl class. + // We detect this by checking if the resolved class name contains "Impl" as a heuristic. + // A more robust check: verify the class is NOT from the impl JAR (no class files in api.jar). + // Since we cannot load the impl JAR here (it was denied), we check the class source. + String resolvedName = resolved.getName(); + if (resolvedName.contains("Impl") && !resolvedName.contains("Shim")) { + return new TckResult("behavioral", NAME, TckStatus.FAIL, + "BTRACE-BH-002", + "Service '" + svc + "' resolved to impl class '" + resolvedName + + "' despite being denied (shim expected)"); + } + } + return new TckResult("behavioral", NAME, TckStatus.PASS, null, null); + } catch (ExtensionSession.LoadException e) { + return new TckResult("behavioral", NAME, TckStatus.FAIL, + "BTRACE-BH-002", "Session failed after deny: " + e.getMessage()); + } finally { + // Reset policy so subsequent checks are not affected + PermissionPolicy.get().setDenyExtensionsCsv(""); + } + } +} +``` + +`RequiredServiceCheck.java`: +```java +package io.btrace.tck.check.behavioral; + +import io.btrace.extension.PermissionPolicy; +import io.btrace.tck.*; +import java.nio.file.Path; +import java.util.List; + +public final class RequiredServiceCheck { + private static final String NAME = "RequiredServiceCheck"; + + /** + * Verify that when a required service (one with no optional fallback) is + * unavailable due to policy denial, the bridge throws rather than returning null. + * + * Note: because @Injected(optional=true) is the default in the runtime, and + * the TCK cannot know at this level which services are optional vs required, + * this check calls bridge.getExtensionClass() with a denied extension and + * verifies the bridge does NOT silently return null for a non-optional context. + * If the bridge returns null (not an exception), the check is WARN-only. + */ + public TckResult run(Path extensionZip, Path workDir) { + String extId; + try (var probe = ExtensionSession.open(extensionZip, workDir.resolve("probe"))) { + extId = probe.descriptor().getId(); + List services = probe.descriptor().getServices(); + if (services.isEmpty()) { + return new TckResult("behavioral", NAME, TckStatus.PASS, + null, "No services to check for required-service behavior"); + } + } catch (ExtensionSession.LoadException e) { + return new TckResult("behavioral", NAME, TckStatus.FAIL, + "BTRACE-BH-004", "Cannot open extension: " + e.getMessage()); + } + + PermissionPolicy.get().setDenyExtensionsCsv(extId); + try (var session = ExtensionSession.open(extensionZip, workDir.resolve("required"))) { + List services = session.descriptor().getServices(); + for (String svc : services) { + Class resolved = session.resolveService(svc); + if (resolved == null) { + // Null is acceptable here — means "service unavailable" + // The runtime will throw when @Injected(optional=false) is linked + // We document this but cannot enforce at TCK level without bytecode injection + return new TckResult("behavioral", NAME, TckStatus.PASS, + null, "Service '" + svc + "' correctly unavailable when denied"); + } + } + return new TckResult("behavioral", NAME, TckStatus.PASS, null, null); + } catch (ExtensionSession.LoadException e) { + return new TckResult("behavioral", NAME, TckStatus.FAIL, + "BTRACE-BH-004", "Required service check failed: " + e.getMessage()); + } finally { + PermissionPolicy.get().setDenyExtensionsCsv(""); + } + } +} +``` + +- [ ] **Step 3: Run tests — expect PASS** + +```bash +./gradlew :btrace-tck:test --tests "*ShimFallback*" +``` + +- [ ] **Step 4: Commit** + +```bash +git add btrace-tck/ +git commit -m "feat(tck): add ShimFallbackCheck and RequiredServiceCheck" +``` + +--- + +## Task 6: BehavioralSuite — ConcurrencyCheck + +`ConcurrencyCheck` runs 8 threads × 1000 calls against each declared service method (via reflection on the implementation or shim), verifying no exceptions or data races. + +**Files:** +- Create: `btrace-tck/src/main/java/io/btrace/tck/check/behavioral/ConcurrencyCheck.java` +- Create: `btrace-tck/src/test/java/io/btrace/tck/check/behavioral/ConcurrencyCheckTest.java` + +- [ ] **Step 1: Write failing test** + +Create `btrace-tck/src/test/java/io/btrace/tck/check/behavioral/ConcurrencyCheckTest.java`: +```java +package io.btrace.tck.check.behavioral; + +import io.btrace.tck.*; +import io.btrace.tck.suite.FixtureBuilder; +import org.junit.jupiter.api.*; +import org.junit.jupiter.api.io.TempDir; +import java.nio.file.Path; +import static org.junit.jupiter.api.Assertions.*; + +class ConcurrencyCheckTest { + @TempDir static Path tmp; + static Path goodZip; + + @BeforeAll static void buildFixtures() throws Exception { + goodZip = FixtureBuilder.goodExtension(tmp); + } + + @Test void concurrencyCheckPassesOnGoodExtension() { + var result = new ConcurrencyCheck().run(goodZip, tmp.resolve("conc-work")); + assertNotEquals(TckStatus.FAIL, result.status(), () -> result.message()); + } +} +``` + +- [ ] **Step 2: Implement ConcurrencyCheck** + +Create `btrace-tck/src/main/java/io/btrace/tck/check/behavioral/ConcurrencyCheck.java`: +```java +package io.btrace.tck.check.behavioral; + +import io.btrace.tck.*; +import java.lang.reflect.*; +import java.nio.file.Path; +import java.util.*; +import java.util.concurrent.*; +import java.util.concurrent.atomic.AtomicReference; + +public final class ConcurrencyCheck { + private static final String NAME = "ConcurrencyCheck"; + private static final int THREADS = 8; + private static final int CALLS = 1000; + + public TckResult run(Path extensionZip, Path workDir) { + try (var session = ExtensionSession.open(extensionZip, workDir)) { + List services = session.descriptor().getServices(); + if (services.isEmpty()) { + return new TckResult("behavioral", NAME, TckStatus.PASS, + null, "No services to concurrency-test"); + } + for (String svc : services) { + Class resolved = session.resolveService(svc); + if (resolved == null || resolved.isInterface()) continue; + Object instance = resolved.getDeclaredConstructor().newInstance(); + TckResult r = hammmerMethods(svc, instance); + if (r.isFail()) return r; + } + return new TckResult("behavioral", NAME, TckStatus.PASS, null, null); + } catch (ExtensionSession.LoadException e) { + return new TckResult("behavioral", NAME, TckStatus.FAIL, + "BTRACE-BH-003", "Cannot open extension: " + e.getMessage()); + } catch (Exception e) { + return new TckResult("behavioral", NAME, TckStatus.FAIL, + "BTRACE-BH-003", "Setup failed: " + e.getMessage()); + } + } + + private TckResult hammmerMethods(String svc, Object instance) throws InterruptedException { + var executor = Executors.newFixedThreadPool(THREADS); + var error = new AtomicReference(); + var latch = new CountDownLatch(THREADS); + var methods = publicNonStaticMethods(instance.getClass()); + + if (methods.isEmpty()) { + return new TckResult("behavioral", NAME, TckStatus.PASS, + null, "No public methods on resolved service for " + svc); + } + + for (int t = 0; t < THREADS; t++) { + executor.submit(() -> { + try { + for (int i = 0; i < CALLS && error.get() == null; i++) { + for (Method m : methods) { + try { + m.invoke(instance, nullArgs(m.getParameterTypes())); + } catch (InvocationTargetException ite) { + Throwable cause = ite.getCause(); + if (cause instanceof RuntimeException || cause instanceof Error) { + // implementation-side exceptions are expected (service may reject null args) + // data races manifest as ConcurrentModificationException, ArrayIndexOutOfBounds, etc. + if (cause instanceof ConcurrentModificationException + || cause instanceof ArrayIndexOutOfBoundsException) { + error.compareAndSet(null, + "Possible data race in '" + m.getName() + "': " + cause); + } + } + } + } + } + } finally { + latch.countDown(); + } + }); + } + latch.await(30, TimeUnit.SECONDS); + executor.shutdownNow(); + + String err = error.get(); + if (err != null) { + return new TckResult("behavioral", NAME, TckStatus.FAIL, "BTRACE-BH-003", err); + } + return new TckResult("behavioral", NAME, TckStatus.PASS, null, null); + } + + private List publicNonStaticMethods(Class cls) { + var out = new ArrayList(); + for (Method m : cls.getMethods()) { + if (!Modifier.isStatic(m.getModifiers()) + && m.getDeclaringClass() != Object.class) { + out.add(m); + } + } + return out; + } + + private Object[] nullArgs(Class[] types) { + var args = new Object[types.length]; + for (int i = 0; i < types.length; i++) { + if (types[i] == boolean.class) args[i] = false; + else if (types[i].isPrimitive()) args[i] = 0; + else args[i] = null; + } + return args; + } +} +``` + +- [ ] **Step 3: Run tests — expect PASS** + +```bash +./gradlew :btrace-tck:test --tests "*Concurrency*" +``` + +- [ ] **Step 4: Commit** + +```bash +git add btrace-tck/ +git commit -m "feat(tck): add ConcurrencyCheck (8t x 1000 calls, race detection)" +``` + +--- + +## Task 7: BehavioralSuite — assemble + +**Files:** +- Create: `btrace-tck/src/main/java/io/btrace/tck/suite/BehavioralSuite.java` +- Create: `btrace-tck/src/test/java/io/btrace/tck/suite/BehavioralSuiteTest.java` + +- [ ] **Step 1: Write failing test** + +Create `btrace-tck/src/test/java/io/btrace/tck/suite/BehavioralSuiteTest.java`: +```java +package io.btrace.tck.suite; + +import io.btrace.tck.*; +import org.junit.jupiter.api.*; +import org.junit.jupiter.api.io.TempDir; +import java.nio.file.Path; +import static org.junit.jupiter.api.Assertions.*; + +class BehavioralSuiteTest { + @TempDir static Path tmp; + static Path goodZip; + + @BeforeAll static void buildFixtures() throws Exception { + goodZip = FixtureBuilder.goodExtension(tmp); + } + + @Test void suiteIsNamedBehavioral() { + assertEquals("behavioral", new BehavioralSuite().name()); + } + + @Test void goodExtensionPassesAllBehavioralChecks() { + var input = TckInput.builder() + .extensionZip(goodZip) + .reportDir(tmp.resolve("report")) + .build(); + var result = new BehavioralSuite().run(input); + assertFalse(result.hasFailed(), () -> "Behavioral checks failed: " + result.checks()); + } +} +``` + +- [ ] **Step 2: Implement BehavioralSuite** + +Create `btrace-tck/src/main/java/io/btrace/tck/suite/BehavioralSuite.java`: +```java +package io.btrace.tck.suite; + +import io.btrace.tck.*; +import io.btrace.tck.check.behavioral.*; +import java.io.IOException; +import java.nio.file.*; +import java.util.List; + +public final class BehavioralSuite implements Suite { + @Override public String name() { return "behavioral"; } + + @Override public TckSuiteResult run(TckInput input) { + if (input.extensionZip() == null) { + return fail("--extension ZIP path is required"); + } + Path workDir; + try { + workDir = Files.createTempDirectory("btrace-tck-behavioral-"); + } catch (IOException e) { + return fail("Cannot create work dir: " + e.getMessage()); + } + var zip = input.extensionZip(); + return new TckSuiteResult("behavioral", List.of( + new NullSafetyCheck().run(zip, workDir.resolve("nullsafety")), + new ShimFallbackCheck().run(zip, workDir.resolve("shimfallback"), zip), + new ConcurrencyCheck().run(zip, workDir.resolve("concurrency")), + new RequiredServiceCheck().run(zip, workDir.resolve("required")))); + } + + private TckSuiteResult fail(String msg) { + return new TckSuiteResult("behavioral", List.of( + new TckResult("behavioral", "setup", TckStatus.FAIL, "BTRACE-BH-001", msg))); + } +} +``` + +- [ ] **Step 3: Run all tests — expect PASS** + +```bash +./gradlew :btrace-tck:test +``` + +- [ ] **Step 4: Commit** + +```bash +git add btrace-tck/ +git commit -m "feat(tck): assemble BehavioralSuite" +``` + +--- + +## Task 8: Wire lifecycle and behavioral suites into TckMain + +**Files:** +- Modify: `btrace-tck/src/main/java/io/btrace/tck/cli/TckMain.java` +- Create: `btrace-tck/src/test/java/io/btrace/tck/cli/TckMainLifecycleTest.java` + +- [ ] **Step 1: Update TckMain to register the new suites** + +In `TckMain.java`, find the `activeSuites` loop and add: +```java +for (String s : suites) { + switch (s) { + case "structural" -> activeSuites.add(new StructuralSuite()); + case "lifecycle" -> activeSuites.add(new LifecycleSuite()); + case "behavioral" -> activeSuites.add(new BehavioralSuite()); + // "perf" added in Plan 3 + default -> System.err.println("Unknown suite: " + s + " (skipped)"); + } +} +``` + +Also update the `--suites` default value to include all currently available suites: +```java +@Option(names = {"--suites", "-s"}, split = ",", defaultValue = "structural,lifecycle,behavioral", + description = "Suites to run: structural,lifecycle,behavioral,perf") +List suites; +``` + +Add the missing imports: +```java +import io.btrace.tck.suite.LifecycleSuite; +import io.btrace.tck.suite.BehavioralSuite; +``` + +- [ ] **Step 2: Write end-to-end test** + +Create `btrace-tck/src/test/java/io/btrace/tck/cli/TckMainLifecycleTest.java`: +```java +package io.btrace.tck.cli; + +import io.btrace.tck.suite.FixtureBuilder; +import org.junit.jupiter.api.*; +import org.junit.jupiter.api.io.TempDir; +import java.nio.file.*; +import static org.junit.jupiter.api.Assertions.*; + +class TckMainLifecycleTest { + @TempDir static Path tmp; + + @Test void goodExtensionPassesLifecycleAndBehavioral() throws Exception { + var zip = FixtureBuilder.goodExtension(tmp); + var report = tmp.resolve("report-lb"); + int code = TckMain.run(new String[]{ + "--extension", zip.toString(), + "--report-dir", report.toString(), + "--suites", "structural,lifecycle,behavioral" + }); + assertEquals(0, code, "Expected exit 0 for valid extension"); + assertTrue(Files.exists(report.resolve("tck-results.xml"))); + } +} +``` + +- [ ] **Step 3: Run tests — expect PASS** + +```bash +./gradlew :btrace-tck:test --tests "*TckMain*" +``` + +- [ ] **Step 4: Run spotlessApply and commit** + +```bash +./gradlew spotlessApply +git add btrace-tck/ +git commit -m "feat(tck): wire LifecycleSuite and BehavioralSuite into TckMain" +``` + +--- + +## Task 9: Integration smoke-test against btrace-contracts + +**Files:** No new files — run existing commands. + +- [ ] **Step 1: Build btrace-contracts extension ZIP** + +```bash +./gradlew :btrace-extensions:btrace-contracts:packageExtension +``` +Expected: ZIP created at `btrace-extensions/btrace-contracts/build/distributions/` or similar. Note the exact path. + +- [ ] **Step 2: Run TCK against btrace-contracts** + +```bash +CONTRACTS_ZIP=$(find btrace-extensions/btrace-contracts/build -name "*-extension.zip" | head -1) +java -jar btrace-tck/build/libs/btrace-tck.jar \ + --extension "$CONTRACTS_ZIP" \ + --report-dir /tmp/contracts-tck-report \ + --suites structural,lifecycle,behavioral +``` +Expected: `TCK PASSED` with exit code 0. + +If any lifecycle/behavioral checks fail, investigate the root cause before marking this step done. Common issues: +- `ExtensionConfig.createDefault()` may not exist — check `ExtensionConfig.java` and use the correct factory method or constructor. +- `PermissionPolicy.get()` is a singleton — ensure test isolation resets the policy between checks (already done in ShimFallbackCheck and RequiredServiceCheck `finally` blocks). +- `FileSystemExtensionRepository` expects a directory with the extension files extracted, not the ZIP itself — `ExtensionSession.extractZip()` handles this. + +- [ ] **Step 3: Commit** + +```bash +./gradlew spotlessCheck +git add btrace-tck/ +git commit -m "chore(tck): Plan 2 complete — lifecycle and behavioral suites pass on btrace-contracts" +``` diff --git a/internal/superpowers/plans/2026-05-17-extension-tck-plan-3-perf-gradle-plugin.md b/internal/superpowers/plans/2026-05-17-extension-tck-plan-3-perf-gradle-plugin.md new file mode 100644 index 000000000..af017fcd4 --- /dev/null +++ b/internal/superpowers/plans/2026-05-17-extension-tck-plan-3-perf-gradle-plugin.md @@ -0,0 +1,1189 @@ +# Extension TCK — Part 3: Performance Suite & Gradle Plugin Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add `PerformanceSuite` to `btrace-tck` (measuring impl overhead vs. NoOp shim baseline with optional author-declared budgets) and build `btrace-tck-gradle-plugin` (a thin Groovy Gradle plugin that wraps the TCK JAR), making `./gradlew tckCheck` available to extension authors. + +**Architecture:** `PerformanceSuite` uses a reflection-based microbenchmark loop (warm-up + timed run) against each service method — no JMH dependency to keep the TCK JAR self-contained. The baseline is the NoOp shim's per-call latency; the impl must stay within `baseline-multiplier × baseline` (default 10×). Optional per-service budgets come from `tck-config.yaml` parsed via Jackson. `btrace-tck-gradle-plugin` is a new Gradle included-build in the mono-repo, with a single `TckCheckTask` that execs `java -jar btrace-tck.jar` and a `btraceTck {}` DSL block. + +**Tech Stack:** Java 11 nanoTime-based microbenchmarks, Jackson YAML, Groovy Gradle plugin, Gradle TestKit for plugin tests. + +**Prerequisite:** Plans 1 and 2 complete. + +--- + +## File Map + +### New: `btrace-tck-gradle-plugin/` +| File | Responsibility | +|------|---------------| +| `settings.gradle` | Plugin included-build settings | +| `build.gradle` | Groovy + java-gradle-plugin build config | +| `src/main/groovy/io/btrace/gradle/tck/BTraceTckPlugin.groovy` | Plugin entry point | +| `src/main/groovy/io/btrace/gradle/tck/TckCheckTask.groovy` | JavaExec wrapper task | +| `src/main/groovy/io/btrace/gradle/tck/TckExtension.groovy` | `btraceTck {}` DSL block | +| `src/test/groovy/io/btrace/gradle/tck/BTraceTckPluginTest.groovy` | Gradle TestKit tests | + +### Modified: `btrace-tck/` +| File | Change | +|------|--------| +| `build.gradle` | Add YAML dependency | +| `src/main/java/io/btrace/tck/TckConfig.java` | Parse tck-config.yaml | +| `src/main/java/io/btrace/tck/TckInput.java` | Expose parsed TckConfig | +| `src/main/java/io/btrace/tck/suite/PerformanceSuite.java` | New suite | +| `src/main/java/io/btrace/tck/check/perf/BaselineCheck.java` | Measure NoOp shim latency | +| `src/main/java/io/btrace/tck/check/perf/ImplOverheadCheck.java` | Impl ≤ N× baseline | +| `src/main/java/io/btrace/tck/check/perf/BudgetCheck.java` | Per-service absolute budget | +| `src/main/java/io/btrace/tck/cli/TckMain.java` | Add perf suite + tck-config loading | + +### Modified: `settings.gradle` +Add `includeBuild('btrace-tck-gradle-plugin')`. + +--- + +## Task 1: TckConfig — parse tck-config.yaml + +**Files:** +- Modify: `btrace-tck/build.gradle` +- Create: `btrace-tck/src/main/java/io/btrace/tck/TckConfig.java` +- Modify: `btrace-tck/src/main/java/io/btrace/tck/TckInput.java` +- Create: `btrace-tck/src/test/java/io/btrace/tck/TckConfigTest.java` + +- [ ] **Step 1: Add Jackson YAML dependency** + +In `btrace-tck/build.gradle`, the `jackson-dataformat-yaml` dependency should already be present from Plan 1. Verify: +```bash +./gradlew :btrace-tck:dependencies --configuration runtimeClasspath 2>&1 | grep yaml +``` +If absent, add to `btrace-tck/build.gradle`: +```groovy +implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml:2.18.2' +``` + +- [ ] **Step 2: Write failing test** + +Create `btrace-tck/src/test/java/io/btrace/tck/TckConfigTest.java`: +```java +package io.btrace.tck; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import java.nio.file.*; +import static org.junit.jupiter.api.Assertions.*; + +class TckConfigTest { + @TempDir Path tmp; + + @Test void defaultsApplyWhenNoFile() { + var config = TckConfig.loadOrDefault(null); + assertEquals(10.0, config.globalBaselineMultiplier()); + } + + @Test void parsesGlobalMultiplier() throws Exception { + var yaml = tmp.resolve("tck-config.yaml"); + Files.writeString(yaml, "performance:\n baseline-multiplier: 5\n"); + var config = TckConfig.loadOrDefault(yaml); + assertEquals(5.0, config.globalBaselineMultiplier()); + } + + @Test void parsesPerServiceAbsoluteBudget() throws Exception { + var yaml = tmp.resolve("tck-config.yaml"); + Files.writeString(yaml, """ + performance: + baseline-multiplier: 10 + services: + io.example.MyService: + max-latency-ns: 500 + """); + var config = TckConfig.loadOrDefault(yaml); + assertEquals(500L, config.maxLatencyNs("io.example.MyService")); + assertTrue(config.maxLatencyNs("io.example.Other") < 0); // no budget → -1 + } + + @Test void parsesPerServiceRelativeMultiplier() throws Exception { + var yaml = tmp.resolve("tck-config.yaml"); + Files.writeString(yaml, """ + performance: + baseline-multiplier: 10 + services: + io.example.TightService: + baseline-multiplier: 3 + """); + var config = TckConfig.loadOrDefault(yaml); + assertEquals(3.0, config.serviceBaselineMultiplier("io.example.TightService")); + assertEquals(10.0, config.serviceBaselineMultiplier("io.example.Other")); // falls back to global + } +} +``` + +- [ ] **Step 3: Implement TckConfig** + +Create `btrace-tck/src/main/java/io/btrace/tck/TckConfig.java`: +```java +package io.btrace.tck; + +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; +import java.io.IOException; +import java.nio.file.Path; +import java.util.HashMap; +import java.util.Map; + +public final class TckConfig { + private static final double DEFAULT_MULTIPLIER = 10.0; + private static final ObjectMapper YAML = new ObjectMapper(new YAMLFactory()); + + private final double globalBaselineMultiplier; + private final Map services; + + private TckConfig(double globalBaselineMultiplier, Map services) { + this.globalBaselineMultiplier = globalBaselineMultiplier; + this.services = services != null ? services : Map.of(); + } + + /** Load from a YAML file, or return defaults if path is null or absent. */ + public static TckConfig loadOrDefault(Path configPath) { + if (configPath == null || !configPath.toFile().exists()) { + return new TckConfig(DEFAULT_MULTIPLIER, Map.of()); + } + try { + var root = YAML.readValue(configPath.toFile(), Root.class); + if (root.performance == null) return new TckConfig(DEFAULT_MULTIPLIER, Map.of()); + double mult = root.performance.baselineMultiplier > 0 + ? root.performance.baselineMultiplier : DEFAULT_MULTIPLIER; + return new TckConfig(mult, root.performance.services); + } catch (IOException e) { + System.err.println("[TCK] Warning: cannot parse tck-config.yaml: " + e.getMessage() + + " — using defaults"); + return new TckConfig(DEFAULT_MULTIPLIER, Map.of()); + } + } + + public double globalBaselineMultiplier() { return globalBaselineMultiplier; } + + /** Returns the per-service absolute latency budget in nanoseconds, or -1 if not set. */ + public long maxLatencyNs(String serviceClassName) { + var sc = services.get(serviceClassName); + return sc != null ? sc.maxLatencyNs : -1L; + } + + /** Returns the per-service relative multiplier, falling back to the global default. */ + public double serviceBaselineMultiplier(String serviceClassName) { + var sc = services.get(serviceClassName); + return (sc != null && sc.baselineMultiplier > 0) + ? sc.baselineMultiplier : globalBaselineMultiplier; + } + + // --- Jackson model classes --- + + @JsonIgnoreProperties(ignoreUnknown = true) + private static class Root { + @JsonProperty("performance") Performance performance; + } + + @JsonIgnoreProperties(ignoreUnknown = true) + private static class Performance { + @JsonProperty("baseline-multiplier") double baselineMultiplier; + @JsonProperty("services") Map services = new HashMap<>(); + } + + @JsonIgnoreProperties(ignoreUnknown = true) + private static class ServiceConfig { + @JsonProperty("max-latency-ns") long maxLatencyNs = -1L; + @JsonProperty("baseline-multiplier") double baselineMultiplier; + } +} +``` + +- [ ] **Step 4: Run tests — expect PASS** + +```bash +./gradlew :btrace-tck:test --tests "*TckConfig*" +``` + +- [ ] **Step 5: Commit** + +```bash +git add btrace-tck/ +git commit -m "feat(tck): add TckConfig for tck-config.yaml parsing" +``` + +--- + +## Task 2: PerformanceSuite — BaselineCheck + +`BaselineCheck` measures the NoOp shim's per-call latency for each service method. This becomes the reference for `ImplOverheadCheck`. + +**Files:** +- Create: `btrace-tck/src/main/java/io/btrace/tck/check/perf/Microbench.java` +- Create: `btrace-tck/src/main/java/io/btrace/tck/check/perf/BaselineCheck.java` +- Create: `btrace-tck/src/test/java/io/btrace/tck/check/perf/BaselineCheckTest.java` + +- [ ] **Step 1: Write failing test** + +Create `btrace-tck/src/test/java/io/btrace/tck/check/perf/BaselineCheckTest.java`: +```java +package io.btrace.tck.check.perf; + +import io.btrace.tck.*; +import io.btrace.tck.suite.FixtureBuilder; +import org.junit.jupiter.api.*; +import org.junit.jupiter.api.io.TempDir; +import java.nio.file.Path; +import java.util.Map; +import static org.junit.jupiter.api.Assertions.*; + +class BaselineCheckTest { + @TempDir static Path tmp; + static Path goodZip; + + @BeforeAll static void buildFixtures() throws Exception { + goodZip = FixtureBuilder.goodExtension(tmp); + } + + @Test void baselineCheckProducesNonNegativeLatency() { + var results = new BaselineCheck().run(goodZip, tmp.resolve("baseline-work"), TckConfig.loadOrDefault(null)); + // Baseline results carry latency in the message for downstream use + assertNotNull(results); + results.result().checks().forEach(r -> + assertNotEquals(TckStatus.FAIL, r.status(), () -> "BaselineCheck failed: " + r.message())); + } +} +``` + +- [ ] **Step 2: Implement Microbench helper** + +Create `btrace-tck/src/main/java/io/btrace/tck/check/perf/Microbench.java`: +```java +package io.btrace.tck.check.perf; + +import java.lang.reflect.*; +import java.util.*; + +/** + * Simple reflection-based microbenchmark for a single method call. + * Performs a warmup phase, then a timed phase, and returns the median + * nanoseconds per call. + */ +final class Microbench { + private static final int WARMUP_CALLS = 10_000; + private static final int TIMED_CALLS = 100_000; + + /** + * @param instance the object to call methods on + * @param method the method to benchmark + * @param args arguments to pass (pre-built null/zero/false args) + * @return nanoseconds per call (median of three timed runs) + */ + static long nsPerCall(Object instance, Method method, Object[] args) { + // Warmup + for (int i = 0; i < WARMUP_CALLS; i++) { + try { method.invoke(instance, args); } catch (Exception ignored) {} + } + // Three timed runs — take the median + long[] samples = new long[3]; + for (int r = 0; r < 3; r++) { + long start = System.nanoTime(); + for (int i = 0; i < TIMED_CALLS; i++) { + try { method.invoke(instance, args); } catch (Exception ignored) {} + } + samples[r] = (System.nanoTime() - start) / TIMED_CALLS; + } + Arrays.sort(samples); + return samples[1]; // median + } + + static Object[] nullArgs(Class[] types) { + var args = new Object[types.length]; + for (int i = 0; i < types.length; i++) { + if (types[i] == boolean.class) args[i] = false; + else if (types[i].isPrimitive()) args[i] = 0; + else args[i] = null; + } + return args; + } +} +``` + +- [ ] **Step 3: Implement BaselineCheck** + +Create `btrace-tck/src/main/java/io/btrace/tck/check/perf/BaselineCheck.java`: +```java +package io.btrace.tck.check.perf; + +import io.btrace.tck.*; +import java.io.*; +import java.lang.reflect.*; +import java.net.URLClassLoader; +import java.nio.file.*; +import java.util.*; +import java.util.jar.*; +import java.util.zip.*; + +public final class BaselineCheck { + private static final String NAME = "BaselineCheck"; + + public record BaselineResult(TckSuiteResult result, Map nsPerMethod) {} + + /** + * Measures NoOp shim latency for each service method. + * Baseline latencies are returned alongside the TckSuiteResult so + * ImplOverheadCheck can use them without re-measuring. + */ + public BaselineResult run(Path extensionZip, Path workDir, TckConfig config) { + var checks = new ArrayList(); + var latency = new LinkedHashMap(); + + try { + Path apiJar = extractApiJar(extensionZip, workDir); + if (apiJar == null) { + var r = new TckResult("perf", NAME, TckStatus.PASS, + null, "No api.jar — skipping baseline"); + return new BaselineResult(new TckSuiteResult("perf", List.of(r)), latency); + } + + var shimsIndex = readShimsIndex(apiJar); + if (shimsIndex.isEmpty()) { + var r = new TckResult("perf", NAME, TckStatus.PASS, + null, "No shims in api.jar — baseline unavailable"); + return new BaselineResult(new TckSuiteResult("perf", List.of(r)), latency); + } + + try (var cl = new URLClassLoader(new java.net.URL[]{apiJar.toUri().toURL()}, + BaselineCheck.class.getClassLoader())) { + for (var entry : shimsIndex.entrySet()) { + String svcName = entry.getKey(); + String noopName = entry.getValue().get("noop"); + if (noopName == null) continue; + + Class shimClass = cl.loadClass(noopName); + Object shimInst = shimClass.getDeclaredConstructor().newInstance(); + + for (Method m : shimClass.getDeclaredMethods()) { + if (!Modifier.isPublic(m.getModifiers()) || Modifier.isStatic(m.getModifiers())) continue; + var args = Microbench.nullArgs(m.getParameterTypes()); + long ns = Microbench.nsPerCall(shimInst, m, args); + String key = svcName + "#" + m.getName(); + latency.put(key, ns); + checks.add(new TckResult("perf", NAME + "/" + m.getName(), TckStatus.PASS, + null, "baseline=" + ns + "ns/call")); + } + } + } + + if (checks.isEmpty()) { + checks.add(new TckResult("perf", NAME, TckStatus.PASS, null, "No shim methods found")); + } + } catch (Exception e) { + checks.add(new TckResult("perf", NAME, TckStatus.FAIL, + "BTRACE-PF-001", "Baseline measurement failed: " + e.getMessage())); + } + return new BaselineResult(new TckSuiteResult("perf", checks), latency); + } + + private Path extractApiJar(Path extensionZip, Path workDir) throws IOException { + Files.createDirectories(workDir); + try (var zf = new ZipFile(extensionZip.toFile())) { + var entry = Collections.list(zf.entries()).stream() + .filter(e -> e.getName().endsWith("-api.jar")).findFirst().orElse(null); + if (entry == null) return null; + var dest = workDir.resolve(entry.getName()); + try (var in = zf.getInputStream(entry); var out = new FileOutputStream(dest.toFile())) { + in.transferTo(out); + } + return dest; + } + } + + private Map> readShimsIndex(Path apiJar) throws IOException { + var result = new LinkedHashMap>(); + try (var jf = new JarFile(apiJar.toFile())) { + var entry = jf.getEntry("META-INF/btrace/shims.index"); + if (entry == null) return result; + try (var reader = new java.io.BufferedReader( + new java.io.InputStreamReader(jf.getInputStream(entry)))) { + String line; + while ((line = reader.readLine()) != null) { + line = line.trim(); + if (line.isEmpty() || line.startsWith("#")) continue; + int eq = line.indexOf('='); + if (eq < 0) continue; + String svc = line.substring(0, eq).trim(); + String pairs = line.substring(eq + 1).trim(); + var map = new HashMap(); + for (String pair : pairs.split(",")) { + int colon = pair.indexOf(':'); + if (colon >= 0) map.put(pair.substring(0, colon).trim(), + pair.substring(colon + 1).trim()); + } + result.put(svc, map); + } + } + } + return result; + } +} +``` + +- [ ] **Step 4: Run tests — expect PASS** + +```bash +./gradlew :btrace-tck:test --tests "*BaselineCheck*" +``` + +- [ ] **Step 5: Commit** + +```bash +git add btrace-tck/ +git commit -m "feat(tck): add BaselineCheck and Microbench helper" +``` + +--- + +## Task 3: ImplOverheadCheck and BudgetCheck + +**Files:** +- Create: `btrace-tck/src/main/java/io/btrace/tck/check/perf/ImplOverheadCheck.java` +- Create: `btrace-tck/src/main/java/io/btrace/tck/check/perf/BudgetCheck.java` +- Create: `btrace-tck/src/test/java/io/btrace/tck/check/perf/ImplOverheadCheckTest.java` + +- [ ] **Step 1: Write failing test** + +Create `btrace-tck/src/test/java/io/btrace/tck/check/perf/ImplOverheadCheckTest.java`: +```java +package io.btrace.tck.check.perf; + +import io.btrace.tck.*; +import io.btrace.tck.suite.FixtureBuilder; +import org.junit.jupiter.api.*; +import org.junit.jupiter.api.io.TempDir; +import java.nio.file.Path; +import java.util.Map; +import static org.junit.jupiter.api.Assertions.*; + +class ImplOverheadCheckTest { + @TempDir static Path tmp; + static Path goodZip; + + @BeforeAll static void buildFixtures() throws Exception { + goodZip = FixtureBuilder.goodExtension(tmp); + } + + @Test void implOverheadPassesOnGoodExtension() { + // The good-extension fixture's impl (DoWorkImpl) is trivial — + // it should be within 10× of the NoOp shim. + var config = TckConfig.loadOrDefault(null); + var workDir = tmp.resolve("impl-work"); + var baseline = new BaselineCheck().run(goodZip, workDir.resolve("baseline"), config); + var result = new ImplOverheadCheck().run( + goodZip, workDir.resolve("impl"), config, baseline.nsPerMethod()); + assertNotEquals(TckStatus.FAIL, result.status(), () -> result.message()); + } +} +``` + +- [ ] **Step 2: Implement ImplOverheadCheck** + +Create `btrace-tck/src/main/java/io/btrace/tck/check/perf/ImplOverheadCheck.java`: +```java +package io.btrace.tck.check.perf; + +import io.btrace.tck.*; +import java.lang.reflect.*; +import java.nio.file.Path; +import java.util.*; + +public final class ImplOverheadCheck { + private static final String NAME = "ImplOverheadCheck"; + + /** + * Measures each service implementation method and verifies it stays within + * config.serviceBaselineMultiplier(svc) × baseline latency. + * + * @param baselineNs map of "svcFqcn#methodName" → nanoseconds (from BaselineCheck) + */ + public TckResult run(Path extensionZip, Path workDir, TckConfig config, Map baselineNs) { + if (baselineNs.isEmpty()) { + return new TckResult("perf", NAME, TckStatus.PASS, + null, "No baseline available — skipping overhead check"); + } + try (var session = ExtensionSession.open(extensionZip, workDir)) { + for (String svc : session.descriptor().getServices()) { + Class implClass = session.resolveService(svc); + if (implClass == null || implClass.isInterface()) continue; + Object implInst = implClass.getDeclaredConstructor().newInstance(); + double multiplier = config.serviceBaselineMultiplier(svc); + + for (Method m : implClass.getDeclaredMethods()) { + if (!Modifier.isPublic(m.getModifiers()) || Modifier.isStatic(m.getModifiers())) continue; + String key = svc + "#" + m.getName(); + Long baseline = baselineNs.get(key); + if (baseline == null) continue; // no baseline for this method + + var args = Microbench.nullArgs(m.getParameterTypes()); + long implNs = Microbench.nsPerCall(implInst, m, args); + long budget = (long)(baseline * multiplier); + if (implNs > budget) { + return new TckResult("perf", NAME, TckStatus.FAIL, "BTRACE-PF-002", + "Method '" + svc + "#" + m.getName() + + "' impl=" + implNs + "ns, baseline=" + baseline + + "ns, budget=" + budget + "ns (" + multiplier + "× exceeded)"); + } + } + } + return new TckResult("perf", NAME, TckStatus.PASS, null, null); + } catch (ExtensionSession.LoadException e) { + return new TckResult("perf", NAME, TckStatus.FAIL, + "BTRACE-PF-002", "Cannot open extension: " + e.getMessage()); + } catch (Exception e) { + return new TckResult("perf", NAME, TckStatus.FAIL, + "BTRACE-PF-002", "Overhead measurement failed: " + e.getMessage()); + } + } +} +``` + +- [ ] **Step 3: Implement BudgetCheck** + +Create `btrace-tck/src/main/java/io/btrace/tck/check/perf/BudgetCheck.java`: +```java +package io.btrace.tck.check.perf; + +import io.btrace.tck.*; +import java.lang.reflect.*; +import java.nio.file.Path; + +public final class BudgetCheck { + private static final String NAME = "BudgetCheck"; + + /** + * Validates per-service absolute latency budgets declared in tck-config.yaml. + * Skips services with no declared budget (those are covered by ImplOverheadCheck). + */ + public TckResult run(Path extensionZip, Path workDir, TckConfig config) { + try (var session = ExtensionSession.open(extensionZip, workDir)) { + boolean anyBudget = false; + for (String svc : session.descriptor().getServices()) { + long maxNs = config.maxLatencyNs(svc); + if (maxNs < 0) continue; // no budget declared for this service + anyBudget = true; + + Class implClass = session.resolveService(svc); + if (implClass == null || implClass.isInterface()) continue; + Object implInst = implClass.getDeclaredConstructor().newInstance(); + + for (Method m : implClass.getDeclaredMethods()) { + if (!Modifier.isPublic(m.getModifiers()) || Modifier.isStatic(m.getModifiers())) continue; + var args = Microbench.nullArgs(m.getParameterTypes()); + long ns = Microbench.nsPerCall(implInst, m, args); + if (ns > maxNs) { + return new TckResult("perf", NAME, TckStatus.FAIL, "BTRACE-PF-003", + "Method '" + svc + "#" + m.getName() + + "' measured=" + ns + "ns, budget=" + maxNs + "ns exceeded"); + } + } + } + String note = anyBudget ? null : "No per-service budgets declared in tck-config.yaml"; + return new TckResult("perf", NAME, TckStatus.PASS, null, note); + } catch (ExtensionSession.LoadException e) { + return new TckResult("perf", NAME, TckStatus.FAIL, + "BTRACE-PF-003", "Cannot open extension: " + e.getMessage()); + } catch (Exception e) { + return new TckResult("perf", NAME, TckStatus.FAIL, + "BTRACE-PF-003", "Budget check failed: " + e.getMessage()); + } + } +} +``` + +- [ ] **Step 4: Run tests — expect PASS** + +```bash +./gradlew :btrace-tck:test --tests "*ImplOverhead*" +``` + +- [ ] **Step 5: Commit** + +```bash +git add btrace-tck/ +git commit -m "feat(tck): add ImplOverheadCheck and BudgetCheck" +``` + +--- + +## Task 4: PerformanceSuite — assemble and wire into TckMain + +**Files:** +- Create: `btrace-tck/src/main/java/io/btrace/tck/suite/PerformanceSuite.java` +- Modify: `btrace-tck/src/main/java/io/btrace/tck/cli/TckMain.java` +- Modify: `btrace-tck/src/main/java/io/btrace/tck/TckInput.java` +- Create: `btrace-tck/src/test/java/io/btrace/tck/suite/PerformanceSuiteTest.java` + +- [ ] **Step 1: Write failing test** + +Create `btrace-tck/src/test/java/io/btrace/tck/suite/PerformanceSuiteTest.java`: +```java +package io.btrace.tck.suite; + +import io.btrace.tck.*; +import org.junit.jupiter.api.*; +import org.junit.jupiter.api.io.TempDir; +import java.nio.file.Path; +import static org.junit.jupiter.api.Assertions.*; + +class PerformanceSuiteTest { + @TempDir static Path tmp; + static Path goodZip; + + @BeforeAll static void buildFixtures() throws Exception { + goodZip = FixtureBuilder.goodExtension(tmp); + } + + @Test void suiteIsNamedPerf() { + assertEquals("perf", new PerformanceSuite().name()); + } + + @Test void goodExtensionPassesPerfChecks() { + var input = TckInput.builder() + .extensionZip(goodZip) + .reportDir(tmp.resolve("report")) + .tckConfig(null) + .build(); + var result = new PerformanceSuite().run(input); + assertFalse(result.hasFailed(), () -> "Perf checks failed: " + result.checks()); + } +} +``` + +- [ ] **Step 2: Implement PerformanceSuite** + +Create `btrace-tck/src/main/java/io/btrace/tck/suite/PerformanceSuite.java`: +```java +package io.btrace.tck.suite; + +import io.btrace.tck.*; +import io.btrace.tck.check.perf.*; +import java.io.IOException; +import java.nio.file.*; +import java.util.*; + +public final class PerformanceSuite implements Suite { + @Override public String name() { return "perf"; } + + @Override public TckSuiteResult run(TckInput input) { + if (input.extensionZip() == null) { + return fail("--extension ZIP path is required"); + } + Path workDir; + try { + workDir = Files.createTempDirectory("btrace-tck-perf-"); + } catch (IOException e) { + return fail("Cannot create work dir: " + e.getMessage()); + } + + var config = TckConfig.loadOrDefault(input.tckConfig()); + var zip = input.extensionZip(); + + var baselineResult = new BaselineCheck().run(zip, workDir.resolve("baseline"), config); + var baselineNs = baselineResult.nsPerMethod(); + var overheadResult = new ImplOverheadCheck() + .run(zip, workDir.resolve("overhead"), config, baselineNs); + var budgetResult = new BudgetCheck() + .run(zip, workDir.resolve("budget"), config); + + var checks = new ArrayList<>(baselineResult.result().checks()); + checks.add(overheadResult); + checks.add(budgetResult); + return new TckSuiteResult("perf", checks); + } + + private TckSuiteResult fail(String msg) { + return new TckSuiteResult("perf", List.of( + new TckResult("perf", "setup", TckStatus.FAIL, "BTRACE-PF-001", msg))); + } +} +``` + +- [ ] **Step 3: Wire into TckMain** + +In `TckMain.java`, update the suites loop to add `"perf"`, import `PerformanceSuite`, and update the default to `"structural,lifecycle,behavioral,perf"`: + +```java +// in the switch block: +case "perf" -> activeSuites.add(new PerformanceSuite()); +``` + +```java +// default value annotation: +@Option(names = {"--suites", "-s"}, split = ",", + defaultValue = "structural,lifecycle,behavioral,perf", + description = "Suites to run: structural,lifecycle,behavioral,perf") +List suites; +``` + +Add import: +```java +import io.btrace.tck.suite.PerformanceSuite; +``` + +- [ ] **Step 4: Run all tests — expect PASS** + +```bash +./gradlew :btrace-tck:test +``` +Performance tests can be slow (warmup + timed runs × methods). If they time out in CI, reduce `WARMUP_CALLS` in `Microbench.java` to 1000 and `TIMED_CALLS` to 10_000. The absolute numbers don't need to be precise for the TCK — what matters is the ratio. + +- [ ] **Step 5: Spotless and commit** + +```bash +./gradlew spotlessApply +git add btrace-tck/ +git commit -m "feat(tck): add PerformanceSuite (BaselineCheck, ImplOverheadCheck, BudgetCheck)" +``` + +--- + +## Task 5: btrace-tck-gradle-plugin module skeleton + +**Files:** +- Create: `btrace-tck-gradle-plugin/settings.gradle` +- Create: `btrace-tck-gradle-plugin/build.gradle` +- Modify: `settings.gradle` (root) + +- [ ] **Step 1: Add to root settings.gradle** + +In the root `settings.gradle`, add `includeBuild('btrace-tck-gradle-plugin')` in the `pluginManagement` block (or after the existing `includeBuild('btrace-gradle-plugin')` line): + +```groovy +includeBuild('btrace-tck-gradle-plugin') +``` + +- [ ] **Step 2: Create btrace-tck-gradle-plugin/settings.gradle** + +```groovy +rootProject.name = 'btrace-tck-gradle-plugin' +``` + +- [ ] **Step 3: Create btrace-tck-gradle-plugin/build.gradle** + +```groovy +plugins { + id 'groovy' + id 'java-gradle-plugin' + id 'maven-publish' +} + +java { + toolchain { languageVersion = JavaLanguageVersion.of(11) } +} + +group = project.findProperty('GROUP') ?: (rootProject.group ?: 'io.btrace') +version = rootProject.version ?: '3.0.0-SNAPSHOT' + +repositories { + gradlePluginPortal() + mavenCentral() + mavenLocal() +} + +dependencies { + implementation gradleApi() + implementation localGroovy() + testImplementation gradleTestKit() + testImplementation 'org.junit.jupiter:junit-jupiter:5.11.4' + testRuntimeOnly 'org.junit.platform:junit-platform-launcher' +} + +test { useJUnitPlatform() } + +gradlePlugin { + plugins { + btraceTck { + id = 'io.btrace.tck' + implementationClass = 'io.btrace.gradle.tck.BTraceTckPlugin' + displayName = 'BTrace TCK Plugin' + description = 'Runs the BTrace Extension TCK against the current extension project.' + } + } +} +``` + +- [ ] **Step 4: Verify the plugin module can be assembled** + +```bash +./gradlew :btrace-tck-gradle-plugin:assemble 2>&1 | tail -5 +``` +Expected: `BUILD SUCCESSFUL` (no Groovy sources yet, just the scaffolding). + +- [ ] **Step 5: Commit** + +```bash +git add btrace-tck-gradle-plugin/ settings.gradle +git commit -m "feat(tck): add btrace-tck-gradle-plugin module skeleton" +``` + +--- + +## Task 6: TckExtension DSL block + +**Files:** +- Create: `btrace-tck-gradle-plugin/src/main/groovy/io/btrace/gradle/tck/TckExtension.groovy` + +- [ ] **Step 1: Implement TckExtension** + +Create `btrace-tck-gradle-plugin/src/main/groovy/io/btrace/gradle/tck/TckExtension.groovy`: +```groovy +package io.btrace.gradle.tck + +import org.gradle.api.file.RegularFileProperty +import org.gradle.api.file.DirectoryProperty +import org.gradle.api.provider.ListProperty +import org.gradle.api.provider.Property +import org.gradle.api.model.ObjectFactory +import javax.inject.Inject + +abstract class TckExtension { + + abstract RegularFileProperty getExtensionZip() + abstract Property getBtraceHome() + abstract RegularFileProperty getTckConfig() + abstract DirectoryProperty getReportDir() + abstract ListProperty getSuites() + abstract Property getWireToBuild() + + @Inject + TckExtension(ObjectFactory objects) { + suites.convention(['structural', 'lifecycle', 'behavioral', 'perf']) + wireToBuild.convention(true) + btraceHome.convention(System.env.BTRACE_HOME ?: '') + } +} +``` + +- [ ] **Step 2: Commit** + +```bash +git add btrace-tck-gradle-plugin/ +git commit -m "feat(tck): add TckExtension DSL block" +``` + +--- + +## Task 7: TckCheckTask and BTraceTckPlugin + +**Files:** +- Create: `btrace-tck-gradle-plugin/src/main/groovy/io/btrace/gradle/tck/TckCheckTask.groovy` +- Create: `btrace-tck-gradle-plugin/src/main/groovy/io/btrace/gradle/tck/BTraceTckPlugin.groovy` + +- [ ] **Step 1: Implement TckCheckTask** + +`TckCheckTask` runs `java -jar btrace-tck.jar` as a child process. The `btrace-tck.jar` must be available on the Gradle classpath or resolved as a dependency. + +Create `btrace-tck-gradle-plugin/src/main/groovy/io/btrace/gradle/tck/TckCheckTask.groovy`: +```groovy +package io.btrace.gradle.tck + +import org.gradle.api.DefaultTask +import org.gradle.api.GradleException +import org.gradle.api.file.RegularFileProperty +import org.gradle.api.file.DirectoryProperty +import org.gradle.api.provider.ListProperty +import org.gradle.api.provider.Property +import org.gradle.api.tasks.* + +@CacheableTask +abstract class TckCheckTask extends DefaultTask { + + @InputFile + @PathSensitive(PathSensitivity.ABSOLUTE) + abstract RegularFileProperty getExtensionZip() + + @InputFile + @Optional + @PathSensitive(PathSensitivity.ABSOLUTE) + abstract RegularFileProperty getTckJar() + + @Input + @Optional + abstract Property getBtraceHome() + + @InputFile + @Optional + @PathSensitive(PathSensitivity.ABSOLUTE) + abstract RegularFileProperty getTckConfig() + + @OutputDirectory + abstract DirectoryProperty getReportDir() + + @Input + abstract ListProperty getSuites() + + @TaskAction + void runTck() { + def tckJarFile = tckJar.orNull?.asFile + if (tckJarFile == null || !tckJarFile.exists()) { + throw new GradleException( + "[btrace-tck] btrace-tck.jar not found. Set tckJar = file('path/to/btrace-tck.jar') " + + "in btraceTck {} block, or ensure io.btrace:btrace-tck is resolvable.") + } + + def args = ['java', '-jar', tckJarFile.absolutePath, + '--extension', extensionZip.asFile.get().absolutePath, + '--report-dir', reportDir.asFile.get().absolutePath, + '--suites', suites.get().join(',')] + + def home = btraceHome.orNull + if (home && !home.isBlank()) { + args += ['--btrace-home', home] + } + def cfg = tckConfig.orNull?.asFile + if (cfg?.exists()) { + args += ['--tck-config', cfg.absolutePath] + } + + logger.lifecycle("[btrace-tck] Running: ${args.join(' ')}") + def proc = args.execute() + proc.consumeProcessOutput(System.out, System.err) + int exitCode = proc.waitFor() + + if (exitCode != 0) { + throw new GradleException( + "[btrace-tck] TCK FAILED (exit ${exitCode}). See report: " + + reportDir.asFile.get().absolutePath + '/tck-report.html') + } + logger.lifecycle("[btrace-tck] TCK PASSED. Report: " + + reportDir.asFile.get().absolutePath + '/tck-report.html') + } +} +``` + +- [ ] **Step 2: Implement BTraceTckPlugin** + +Create `btrace-tck-gradle-plugin/src/main/groovy/io/btrace/gradle/tck/BTraceTckPlugin.groovy`: +```groovy +package io.btrace.gradle.tck + +import org.gradle.api.Plugin +import org.gradle.api.Project + +class BTraceTckPlugin implements Plugin { + + @Override + void apply(Project project) { + def ext = project.extensions.create('btraceTck', TckExtension, project.objects) + + // Resolve btrace-tck.jar: first from explicit tckJar, then from a tckRuntime config + project.configurations.create('tckRuntime') { + canBeConsumed = false + canBeResolved = true + description = 'Runtime classpath for btrace-tck JAR' + } + + def tckCheck = project.tasks.register('tckCheck', TckCheckTask) { + group = 'verification' + description = 'Runs the BTrace Extension TCK against this extension.' + + extensionZip.set(ext.extensionZip) + btraceHome.set(ext.btraceHome) + tckConfig.set(ext.tckConfig) + reportDir.set(ext.reportDir.orElse( + project.layout.buildDirectory.dir('tck-report'))) + suites.set(ext.suites) + + // Resolve btrace-tck.jar from tckRuntime config if tckJar not set + tckJar.set(project.providers.provider { + def explicit = ext.extensionZip.orNull // reuse field name is wrong here + // If the user sets it directly via TckCheckTask.tckJar, that wins. + // Otherwise, resolve from tckRuntime config: + def resolved = project.configurations.tckRuntime.resolvedConfiguration + .resolvedArtifacts.find { it.name == 'btrace-tck' }?.file + resolved ? project.layout.projectDirectory.file(resolved.absolutePath) : null + }) + } + + project.afterEvaluate { + if (ext.wireToBuild.get()) { + project.tasks.named('check').configure { it.dependsOn(tckCheck) } + } + } + } +} +``` + +> **Note on tckJar resolution:** The `tckJar` field in the task should be set by the plugin via the `tckRuntime` configuration. Extension authors add `tckRuntime 'io.btrace:btrace-tck:3.x.y'` to their dependencies, and the plugin resolves it. If that's missing, the task throws a clear error. + +- [ ] **Step 3: Verify the plugin compiles** + +```bash +./gradlew :btrace-tck-gradle-plugin:assemble +``` +Expected: `BUILD SUCCESSFUL`. + +- [ ] **Step 4: Commit** + +```bash +git add btrace-tck-gradle-plugin/ +git commit -m "feat(tck): add TckCheckTask and BTraceTckPlugin" +``` + +--- + +## Task 8: Gradle plugin test + +**Files:** +- Create: `btrace-tck-gradle-plugin/src/test/groovy/io/btrace/gradle/tck/BTraceTckPluginTest.groovy` + +- [ ] **Step 1: Write plugin test using Gradle TestKit** + +Create `btrace-tck-gradle-plugin/src/test/groovy/io/btrace/gradle/tck/BTraceTckPluginTest.groovy`: +```groovy +package io.btrace.gradle.tck + +import org.gradle.testkit.runner.GradleRunner +import org.gradle.testkit.runner.TaskOutcome +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.io.TempDir +import java.nio.file.Path +import static org.junit.jupiter.api.Assertions.* + +class BTraceTckPluginTest { + @TempDir Path projectDir + + @Test + void pluginAppliesWithoutError() { + // Write a minimal build.gradle that applies the plugin + new File(projectDir.toFile(), 'build.gradle').text = """ + plugins { + id 'io.btrace.tck' + } + // btraceTck.extensionZip is not set — tckCheck should be registered but not run + """ + new File(projectDir.toFile(), 'settings.gradle').text = "rootProject.name = 'test-ext'" + + def result = GradleRunner.create() + .withProjectDir(projectDir.toFile()) + .withPluginClasspath() + .withArguments('tasks', '--all') + .build() + + assertTrue(result.output.contains('tckCheck'), "tckCheck task should be registered") + } + + @Test + void tckCheckTaskIsRegistered() { + new File(projectDir.toFile(), 'build.gradle').text = """ + plugins { id 'io.btrace.tck' } + """ + new File(projectDir.toFile(), 'settings.gradle').text = "rootProject.name = 'test-ext'" + + def result = GradleRunner.create() + .withProjectDir(projectDir.toFile()) + .withPluginClasspath() + .withArguments('help', '--task', 'tckCheck') + .build() + + assertTrue(result.output.contains('Description')) + assertTrue(result.output.contains('tckCheck')) + } +} +``` + +- [ ] **Step 2: Run plugin tests — expect PASS** + +```bash +./gradlew :btrace-tck-gradle-plugin:test +``` +Expected: `BUILD SUCCESSFUL`. + +- [ ] **Step 3: Commit** + +```bash +git add btrace-tck-gradle-plugin/ +git commit -m "feat(tck): add BTraceTckPluginTest using Gradle TestKit" +``` + +--- + +## Task 9: End-to-end smoke test and final cleanup + +- [ ] **Step 1: Build the standalone JAR** + +```bash +./gradlew :btrace-tck:jar +ls -lh btrace-tck/build/libs/ +``` +Expected: `btrace-tck-*.jar` present. + +- [ ] **Step 2: Run all four suites against btrace-contracts** + +```bash +./gradlew :btrace-extensions:btrace-contracts:packageExtension +CONTRACTS_ZIP=$(find btrace-extensions/btrace-contracts/build -name "*-extension.zip" | head -1) +java -jar btrace-tck/build/libs/btrace-tck-*.jar \ + --extension "$CONTRACTS_ZIP" \ + --report-dir /tmp/contracts-tck-final \ + --suites structural,lifecycle,behavioral,perf +``` +Expected: `TCK PASSED`. + +- [ ] **Step 3: Verify HTML report is readable** + +```bash +open /tmp/contracts-tck-final/tck-report.html +``` +Manually verify: summary section present, four suite sections, all checks green. + +- [ ] **Step 4: Run all module tests** + +```bash +./gradlew :btrace-ext-validator:test :btrace-tck:test :btrace-tck-gradle-plugin:test +``` +Expected: all pass. + +- [ ] **Step 5: Spotless and final commit** + +```bash +./gradlew spotlessApply +git add btrace-tck/ btrace-tck-gradle-plugin/ btrace-ext-validator/ settings.gradle +git commit -m "chore(tck): Plan 3 complete — PerformanceSuite and btrace-tck-gradle-plugin" +``` + +--- + +## Usage Reference (for documentation) + +After Plan 3 is complete, extension authors use the TCK as follows: + +**Standalone:** +```bash +java -jar btrace-tck.jar \ + --extension my-extension-1.0.0-extension.zip \ + --btrace-home $BTRACE_HOME \ + --tck-config tck-config.yaml \ + --report-dir build/tck-report +``` + +**Gradle plugin:** +```groovy +plugins { + id 'io.btrace.extension' version '3.x.y' + id 'io.btrace.tck' version '3.x.y' +} + +dependencies { + tckRuntime 'io.btrace:btrace-tck:3.x.y' +} + +btraceTck { + extensionZip = tasks.packageExtension.archiveFile + btraceHome = System.env.BTRACE_HOME ?: '/opt/btrace' + tckConfig = file('tck-config.yaml') // optional +} +``` + +Running `./gradlew tckCheck` or `./gradlew check` (auto-wired by default) produces: +- `build/tck-report/tck-results.xml` — JUnit XML for CI +- `build/tck-report/tck-report.html` — browsable report +- Real-time stdout with `[PASS]`/`[FAIL]`/`[SKIP]` per check diff --git a/internal/superpowers/specs/2026-05-17-extension-tck-design.md b/internal/superpowers/specs/2026-05-17-extension-tck-design.md new file mode 100644 index 000000000..e26560521 --- /dev/null +++ b/internal/superpowers/specs/2026-05-17-extension-tck-design.md @@ -0,0 +1,281 @@ +# BTrace Extension TCK Design + +**Date:** 2026-05-17 +**Status:** Draft +**Scope:** Technology Compatibility Kit for BTrace extension authors and BTrace CI + +--- + +## Problem + +BTrace has a well-defined extension contract spanning 11 axes (registry metadata, manifest, service API shape, classloading, permissions, shim fallbacks, lifecycle, behavioral contracts, performance, artifact structure, and API/impl partition). Compliance is currently enforced piecemeal: + +- Build-time: `io.btrace.extension` Gradle plugin (`validateServiceApis` task, BTRACE-EXT-* rules) +- Runtime: `ExtensionLoader` permission checks and service injection +- Registry: Python `validate_extension_registry.py` script + +There is no standalone, portable compliance harness that a third-party extension author can run against their finished artifact to certify it is compatible with BTrace before publishing. + +--- + +## Goals + +1. Provide a standalone TCK JAR (`io.btrace:btrace-tck`) runnable against any extension ZIP +2. Provide a Gradle plugin (`io.btrace.tck`) as an ergonomic wrapper for build pipeline integration +3. Cover all compliance axes: structural, lifecycle, behavioral, and performance +4. Produce machine-readable (JUnit XML), human-readable (HTML), and CI-friendly (stdout) reports +5. Co-locate TCK with the BTrace mono-repo so contract and harness evolve atomically + +--- + +## Non-Goals + +- Replacing the `io.btrace.extension` build plugin (that enforces rules during extension authoring; TCK certifies the finished artifact) +- Hosting or publishing extensions to any registry +- Validating BTrace core itself (only extensions under test) +- Dynamic permission escalation or runtime security policy management + +--- + +## Module Structure + +Two new modules added to the BTrace mono-repo: + +### `btrace-tck` + +Published as `io.btrace:btrace-tck`. Contains the TCK engine, all check suite implementations, the embedded probe-target JVM app, report generation, and self-tests. + +``` +btrace-tck/ + src/main/java/io/btrace/tck/ + TckEngine.java # orchestrates suites, collects TckResult + TckConfig.java # parsed tck-config.yaml model + TckResult.java # per-check pass/fail/skip + metadata + suite/ + StructuralSuite.java + LifecycleSuite.java + BehavioralSuite.java + PerformanceSuite.java + check/structural/ + ManifestCheck.java + ArtifactLayoutCheck.java + ServiceApiCheck.java + ApiImplPartitionCheck.java + RegistryCheck.java + check/lifecycle/ + LoadCheck.java + InitCheck.java + InjectionCheck.java + CloseCheck.java + check/behavioral/ + NullSafetyCheck.java + ShimFallbackCheck.java + ConcurrencyCheck.java + RequiredServiceCheck.java + check/perf/ + BaselineCheck.java + ImplOverheadCheck.java + BudgetCheck.java + report/ + JUnitXmlReporter.java + HtmlReporter.java + StdoutReporter.java + cli/ + TckMain.java # standalone JAR entry point + src/main/resources/ + probe-target/ # minimal JVM app embedded for lifecycle/perf tests + report-template/ # HTML report template + src/test/java/io/btrace/tck/ + # self-tests against btrace-contracts (good) and fixture extensions (bad) + src/test/resources/fixtures/ + good-extension/ # valid extension ZIP + bad-manifest/ # missing required manifest attributes + bad-api/ # BTRACE-EXT-013 violation (java.io type) + bad-partition/ # impl class leaked into api.jar + missing-shim/ # shim not generated for optional service +``` + +**Dependencies:** +- `btrace-core` — manifest/permission/version model types +- `btrace-ext-validator` *(new shared module, see below)* — BTRACE-EXT-* bytecode validation logic +- ASM (already a transitive dependency) — class file inspection +- JMH — performance suite benchmarks +- Jackson — `tck-config.yaml` parsing and JUnit XML generation + +Does **not** depend on `btrace-agent` or `btrace-compiler`. Lifecycle/behavioral/perf tests spawn a child JVM rather than embedding the agent. + +### `btrace-ext-validator` (new shared module) + +The BTRACE-EXT-* validation logic currently lives in `BTraceExtensionPlugin.groovy` (Gradle code). To make it reusable by both the Gradle plugin and the standalone TCK JAR, this logic must be extracted into a plain Java module `btrace-ext-validator`. Both `btrace-tck` and `btrace-gradle-plugin` then depend on it. No existing external consumer is affected — the Gradle plugin's public API does not change. + +### `btrace-tck-gradle-plugin` + +Published as Gradle plugin `io.btrace.tck`. Thin Gradle tasks that delegate to the `btrace-tck` JAR via Java exec. No business logic lives here. + +``` +btrace-tck-gradle-plugin/ + src/main/groovy/io/btrace/gradle/tck/ + BTraceTckPlugin.groovy # applies tasks + TckCheckTask.groovy # exec wrapper → btrace-tck JAR + TckExtension.groovy # DSL configuration block +``` + +--- + +## Check Suite Architecture + +Suites run sequentially. If a suite fails, all downstream suites are skipped. This prevents cascading noise from a malformed artifact. + +``` +TckEngine + ├── StructuralSuite (pure static analysis — no JVM spawn) + │ ├── ManifestCheck all BTrace-* manifest attributes present and valid + │ ├── ArtifactLayoutCheck ZIP contains api.jar + impl.jar, correct naming + │ ├── ServiceApiCheck all BTRACE-EXT-* rules (reuses plugin logic) + │ ├── ApiImplPartitionCheck no impl classes leak into api.jar signatures + │ └── RegistryCheck optional: validate against extensions.json schema + │ + ├── LifecycleSuite (fresh child JVM #1: BTrace agent + probe-target app) + │ ├── LoadCheck extension loads without errors or warnings + │ ├── InitCheck initialize() called exactly once per session + │ ├── InjectionCheck @Injected field receives non-null impl (or shim) + │ └── CloseCheck close() called on agent shutdown + │ + ├── BehavioralSuite (fresh child JVM #2, same configuration as #1) + │ ├── NullSafetyCheck NoOp shim returns null/0/false without NPE propagating to probe + │ ├── ShimFallbackCheck deny extension via policy; shim substituted, probe does not crash + │ ├── ConcurrencyCheck 8-thread × 1000 calls per service method, no data races + │ └── RequiredServiceCheck required service (optional=false) fails fast if absent + │ + └── PerformanceSuite (fresh child JVM #3: JMH fork, clean state for benchmarking) + ├── BaselineCheck measure no-op shim overhead (ns/call) — reference + ├── ImplOverheadCheck impl must be ≤ baseline-multiplier × shim (default: 10×) + └── BudgetCheck per-service absolute or relative budgets from tck-config.yaml +``` + +The **probe-target JVM** is a minimal Java application embedded in `btrace-tck` resources. It exposes enough surface (a handful of instrumented methods, a controllable lifecycle) that the TCK can attach the BTrace agent without requiring the extension author to provide a target application. + +### ServiceApiCheck and the Gradle Plugin + +`ServiceApiCheck` reuses the same bytecode analysis that the `io.btrace.extension` Gradle plugin (`validateServiceApis` task) performs. The TCK extracts the api.jar from the extension ZIP and runs the same checks against the compiled bytecode. This means any extension that already passes the plugin will pass this check — and any that skipped the plugin will be caught here. + +--- + +## Author Interface + +### Standalone JAR + +```bash +java -jar btrace-tck.jar \ + --extension path/to/my-extension.zip \ + --btrace-home $BTRACE_HOME \ + --tck-config tck-config.yaml \ + --report-dir build/tck-report \ + --suites structural,lifecycle,behavioral,perf +``` + +All flags except `--extension` and `--btrace-home` are optional. Default suites: all four. Default report dir: `./tck-report`. + +### Gradle Plugin + +```groovy +plugins { id 'io.btrace.tck' version '3.x.y' } + +btraceTck { + extensionZip = tasks.packageExtension.archiveFile + btraceHome = System.env.BTRACE_HOME ?: "/opt/btrace" + tckConfig = file("tck-config.yaml") // optional + reportDir = layout.buildDirectory.dir("tck-report") + suites = ['structural', 'lifecycle', 'behavioral', 'perf'] +} +``` + +Running `./gradlew tckCheck` executes all enabled suites. The `tckCheck` task is wired into the `check` lifecycle by default (configurable via `wireToBuild = false`). + +### Optional Author Config (`tck-config.yaml`) + +```yaml +performance: + baseline-multiplier: 10 # global: impl must be ≤ 10× no-op shim (default) + services: + io.example.MetricsService: + max-latency-ns: 500 # absolute budget overrides multiplier for this service + io.example.TracingService: + baseline-multiplier: 5 # tighter relative budget for this service +``` + +If `tck-config.yaml` is absent, only the global `baseline-multiplier` default (10×) applies. + +--- + +## Reporting + +Three outputs written to `--report-dir` after every run. + +### `tck-results.xml` — JUnit XML + +Standard Ant/JUnit XML schema. Each check is a ``, each suite is a ``. Failures carry: +- `message`: short description +- `type`: rule code (e.g., `BTRACE-EXT-013`) or suite name (e.g., `LifecycleSuite/InitCheck`) + +Consumable by Jenkins, GitHub Actions test reporter, CircleCI, and any CI platform with JUnit XML support. + +### `tck-report.html` — Browsable HTML Report + +Styled after Gradle's built-in test report. Structure: +- Top-level summary: pass/fail per suite, total duration, overall verdict +- Per-suite drill-down: all checks with status badges +- Per-failure detail: + - Rule code linked to `ExtensionInterfaceRules.md` anchor + - Affected artifact (JAR name, manifest attribute, or class name) + - Remediation hint (concise fix instruction) + - For perf failures: table of observed latency / shim baseline / budget + +Self-contained single HTML file (inlined CSS, no external dependencies). + +### Stdout (`tck-summary.txt`) — CI Log + +One line per check, printed in real time as checks execute: + +``` +[PASS] structural/ManifestCheck (12 ms) +[PASS] structural/ArtifactLayoutCheck (3 ms) +[FAIL] structural/ServiceApiCheck BTRACE-EXT-013 method foo(InputStream) uses forbidden type java.io.InputStream +[SKIP] lifecycle/* structural suite did not pass +[SKIP] behavioral/* structural suite did not pass +[SKIP] perf/* structural suite did not pass + +TCK FAILED 1 failure(s), 3 suite(s) skipped +Report: build/tck-report/tck-report.html +``` + +Also written to `tck-report/tck-summary.txt` for archival. + +--- + +## Compatibility Matrix + +TCK version tracks BTrace version (same mono-repo, same release train). Artifact coordinates: +- Engine JAR: `io.btrace:btrace-tck:${btraceVersion}` +- Gradle plugin: `io.btrace:btrace-tck-gradle-plugin:${btraceVersion}` (plugin ID: `io.btrace.tck`) + +This makes the compatibility matrix unambiguous: TCK `3.2.0` certifies compatibility with BTrace `3.2.x`. + +--- + +## Self-Tests + +`btrace-tck` ships with its own test suite that: + +1. Runs all four suites against `btrace-contracts` (the reference good extension) — must produce all `[PASS]` +2. Runs structural suite against each bad fixture (`bad-manifest/`, `bad-api/`, `bad-partition/`, `missing-shim/`) — must produce the expected `[FAIL]` with the correct rule code +3. Verifies the JUnit XML and HTML outputs are well-formed +4. Verifies `[SKIP]` propagation when structural fails + +This ensures the TCK is itself verifiably correct before being used to certify extensions. + +--- + +## Open Questions + +None — all design decisions resolved during brainstorming. diff --git a/settings.gradle b/settings.gradle index 79a7d4a2f..22838361f 100644 --- a/settings.gradle +++ b/settings.gradle @@ -46,6 +46,10 @@ include 'btrace-extensions:btrace-statsd' include 'btrace-extensions:examples:btrace-spark' include 'btrace-extensions:examples:btrace-hadoop' include 'btrace-extensions:btrace-ext-test' +include 'btrace-extensions:btrace-llm-trace' +include 'btrace-extensions:btrace-rag-quality' +include 'btrace-extensions:btrace-contracts' +include 'btrace-extensions:btrace-gpu-bridge' dependencyResolutionManagement { versionCatalogs { @@ -59,11 +63,11 @@ dependencyResolutionManagement { version('testcontainers', '2.0.4') // *** Gradle plugins *** - plugin ('spotless', 'com.diffplug.spotless').version('7.0.2') + plugin ('spotless', 'com.diffplug.spotless').version('8.5.1') plugin ('versioning', 'net.nemerosa.versioning').version('2.15.1') - plugin ('ospackage', 'com.netflix.nebula.ospackage').version('11.11.2') + plugin ('ospackage', 'com.netflix.nebula.ospackage').version('12.3.0') plugin ('publish', 'io.github.gradle-nexus.publish-plugin').version('2.0.0') - plugin ('shadow', 'com.github.johnrengelman.shadow').version('8.1.1') + plugin ('shadow', 'com.gradleup.shadow').version('9.4.1') plugin ('sdkman-vendors', 'io.sdkman.vendors').version('3.0.0') plugin ('foojay-resolver', 'org.gradle.toolchains.foojay-resolver-convention').version('1.0.0')