Restore the agent Feature Flags behaviors through the bridge

dd-java-agent/instrumentation/datadog/openfeature/dd-openfeature-1.69/src/main/java/datadog/trace/instrumentation/openfeature/DDOpenFeatureInstrumentation.java

@@ -41,6 +41,7 @@ public String[] helperClassNames() {
       packageName + ".AgentConfigurationSource",
       packageName + ".AgentEventProxy",
       packageName + ".AgentHealthMetrics",
+      packageName + ".AgentShutdownSource",
       packageName + ".AgentSpanEnricher",
       packageName + ".JavaAgentBridge",
     };

dd-java-agent/instrumentation/datadog/openfeature/dd-openfeature-1.69/src/main/java11/datadog/trace/instrumentation/openfeature/AgentShutdownSource.java

@@ -0,0 +1,18 @@
+package datadog.trace.instrumentation.openfeature;
+
+import com.datadog.openfeature.internal.bridge.ShutdownSource;
+import datadog.trace.api.featureflag.FeatureFlaggingGateway;
+
+/** Notifies the agent shutdown, before the JVM shutdown hooks run. */
+public final class AgentShutdownSource implements ShutdownSource {
+  private final FeatureFlaggingGateway.Backend backend;
+
+  AgentShutdownSource(final FeatureFlaggingGateway.Backend backend) {
+    this.backend = backend;
+  }
+
+  @Override
+  public AutoCloseable subscribe(final Runnable listener) {
+    return this.backend.subscribeShutdown(listener);
+  }
+}

dd-java-agent/instrumentation/datadog/openfeature/dd-openfeature-1.69/src/main/java11/datadog/trace/instrumentation/openfeature/JavaAgentBridge.java

@@ -1,16 +1,29 @@
 package datadog.trace.instrumentation.openfeature;
 
+import static java.lang.reflect.Modifier.isAbstract;
+import static java.lang.reflect.Modifier.isStatic;
+
 import com.datadog.openfeature.internal.bridge.AgentBridge;
 import com.datadog.openfeature.internal.bridge.ConfigurationSource;
 import com.datadog.openfeature.internal.bridge.EventTransport;
 import com.datadog.openfeature.internal.bridge.HealthMetrics;
 import com.datadog.openfeature.internal.bridge.ShutdownSource;
 import com.datadog.openfeature.internal.bridge.SpanEnricher;
 import datadog.trace.api.featureflag.FeatureFlaggingGateway;
+import java.lang.reflect.Method;
 import javax.annotation.Nullable;
 
 /** The bridge backed by the agent Feature Flags backend. */
 public final class JavaAgentBridge implements AgentBridge {
+  // Muzzle checks calls, but not the SPI methods implemented by these helpers.
+  private static final boolean SPI_COMPATIBLE =
+      implementsSpi(AgentBridge.class, JavaAgentBridge.class)
+          && implementsSpi(ConfigurationSource.class, AgentConfigurationSource.class)
+          && implementsSpi(EventTransport.class, AgentEventProxy.class)
+          && implementsSpi(HealthMetrics.class, AgentHealthMetrics.class)
+          && implementsSpi(SpanEnricher.class, AgentSpanEnricher.class)
+          && implementsSpi(ShutdownSource.class, AgentShutdownSource.class);
+
   private final FeatureFlaggingGateway.Backend backend;
   private final HealthMetrics healthMetrics;
   private final SpanEnricher spanEnricher;
@@ -24,12 +37,32 @@ public final class JavaAgentBridge implements AgentBridge {
   /**
    * Connects the SDK to the agent backend.
    *
-   * @param fallback the bridge to use if the agent backend is not started.
+   * @param fallback the bridge to use if the agent backend is not started or the SDK SPI is
+   *     incompatible.
    * @return the agent bridge, or the fallback one.
    */
   public static AgentBridge connect(final AgentBridge fallback) {
     final FeatureFlaggingGateway.Backend backend = FeatureFlaggingGateway.backend();
-    return backend == null ? fallback : new JavaAgentBridge(backend);
+    return backend == null || !SPI_COMPATIBLE ? fallback : new JavaAgentBridge(backend);
+  }
+
+  private static boolean implementsSpi(final Class<?> spi, final Class<?> implementation) {
+    try {
+      for (final Method method : spi.getMethods()) {
+        if (!isStatic(method.getModifiers())) {
+          final Method implemented =
+              implementation.getMethod(method.getName(), method.getParameterTypes());
+          if (isAbstract(implemented.getModifiers())
+              || isStatic(implemented.getModifiers())
+              || implemented.getReturnType() != method.getReturnType()) {
+            return false;
+          }
+        }
+      }
+      return true;
+    } catch (final NoSuchMethodException e) {
+      return false;
+    }
   }
 
   @Nullable
@@ -64,7 +97,6 @@ public SpanEnricher spanEnricher() {
 
   @Override
   public ShutdownSource shutdown() {
-    // TODO Flush the SDK through the agent shutdown. Until then, the SDK uses a JVM shutdown hook.
-    return null;
+    return new AgentShutdownSource(this.backend);
   }
 }

dd-java-agent/instrumentation/datadog/openfeature/dd-openfeature-1.69/src/test/java/datadog/trace/instrumentation/openfeature/DDOpenFeatureClassLoaderTest.java

@@ -0,0 +1,197 @@
+package datadog.trace.instrumentation.openfeature;
+
+import static java.nio.file.Files.createDirectories;
+import static java.nio.file.Files.write;
+import static net.bytebuddy.jar.asm.Opcodes.ACC_ABSTRACT;
+import static net.bytebuddy.jar.asm.Opcodes.ACC_PUBLIC;
+import static net.bytebuddy.jar.asm.Opcodes.ACONST_NULL;
+import static net.bytebuddy.jar.asm.Opcodes.ARETURN;
+import static net.bytebuddy.jar.asm.Opcodes.ASM9;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import com.datadog.openfeature.internal.bridge.AgentBridge;
+import datadog.trace.agent.test.AbstractInstrumentationTest;
+import datadog.trace.api.featureflag.FeatureFlaggingGateway;
+import java.io.IOException;
+import java.io.InputStream;
+import java.net.URL;
+import java.net.URLClassLoader;
+import java.nio.file.Path;
+import java.util.function.UnaryOperator;
+import net.bytebuddy.jar.asm.ClassReader;
+import net.bytebuddy.jar.asm.ClassVisitor;
+import net.bytebuddy.jar.asm.ClassWriter;
+import net.bytebuddy.jar.asm.MethodVisitor;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+import org.tabletest.junit.TableTest;
+
+class DDOpenFeatureClassLoaderTest extends AbstractInstrumentationTest {
+  @TempDir Path sdkOverrides;
+
+  @AfterEach
+  void unregisterBackend() {
+    FeatureFlaggingGateway.register(null);
+  }
+
+  @Test
+  void injectsIntoAnIsolatedSdkClassLoader() throws Exception {
+    final FeatureFlaggingGateway.Backend backend = mock(FeatureFlaggingGateway.Backend.class);
+    when(backend.setting("site")).thenReturn("datadoghq.eu");
+    FeatureFlaggingGateway.register(backend);
+
+    try (URLClassLoader loader = sdkClassLoader()) {
+      final Object bridge = detect(loader);
+      assertSame(loader, bridge.getClass().getClassLoader());
+      assertEquals(JavaAgentBridge.class.getName(), bridge.getClass().getName());
+      assertEquals(
+          "datadoghq.eu",
+          loader
+              .loadClass(AgentBridge.class.getName())
+              .getMethod("setting", String.class)
+              .invoke(bridge, "site"));
+    }
+  }
+
+  @TableTest({
+    "Scenario              | SPI                 | Method    | Return type       ",
+    "settings changed      | AgentBridge         | setting   |                   ",
+    "configuration changed | ConfigurationSource | subscribe |                   ",
+    "transport changed     | EventTransport      | post      |                   ",
+    "metrics changed       | HealthMetrics       | count     |                   ",
+    "enrichment changed    | SpanEnricher        | serialId  |                   ",
+    "shutdown changed      | ShutdownSource      | subscribe |                   ",
+    "return type changed   | AgentBridge         | setting   | Ljava/lang/Object;"
+  })
+  void keepsTheStandaloneBridgeWhenAnSpiMethodChanges(
+      final String spi, final String method, final String returnType) throws Exception {
+    rewriteSpi(
+        spi,
+        bytes -> {
+          final ClassWriter writer = new ClassWriter(0);
+          new ClassReader(bytes)
+              .accept(
+                  new ClassVisitor(ASM9, writer) {
+                    @Override
+                    public MethodVisitor visitMethod(
+                        int access,
+                        String name,
+                        String descriptor,
+                        String signature,
+                        String[] exceptions) {
+                      if (name.equals(method)) {
+                        if (returnType == null) {
+                          name = "renamed";
+                        } else {
+                          descriptor =
+                              descriptor.substring(0, descriptor.indexOf(')') + 1) + returnType;
+                        }
+                      }
+                      return super.visitMethod(access, name, descriptor, signature, exceptions);
+                    }
+                  },
+                  0);
+          return writer.toByteArray();
+        });
+    FeatureFlaggingGateway.register(mock(FeatureFlaggingGateway.Backend.class));
+
+    try (URLClassLoader loader = sdkClassLoader()) {
+      // Implemented methods are not calls, so Muzzle alone accepts this incompatible SPI.
+      assertTrue(new DDOpenFeatureInstrumentation().getInstrumentationMuzzle().matches(loader));
+      assertSame(
+          loader.loadClass(AgentBridge.class.getName()).getField("NONE").get(null), detect(loader));
+    }
+  }
+
+  @TableTest({
+    "Scenario                       | Abstract | Static conflict",
+    "new abstract method            | true     | false          ",
+    "new default method             | false    | false          ",
+    "abstract conflicts with static | true     | true           ",
+    "default conflicts with static  | false    | true           "
+  })
+  void checksNewSpiMethods(final boolean isAbstract, final boolean staticConflict)
+      throws Exception {
+    rewriteSpi(
+        "AgentBridge",
+        bytes -> {
+          final ClassWriter writer = new ClassWriter(0);
+          new ClassReader(bytes)
+              .accept(
+                  new ClassVisitor(ASM9, writer) {
+                    @Override
+                    public void visitEnd() {
+                      final MethodVisitor method =
+                          super.visitMethod(
+                              ACC_PUBLIC | (isAbstract ? ACC_ABSTRACT : 0),
+                              staticConflict ? "connect" : "additionalMethod",
+                              staticConflict
+                                  ? "(Lcom/datadog/openfeature/internal/bridge/AgentBridge;)"
+                                      + "Lcom/datadog/openfeature/internal/bridge/AgentBridge;"
+                                  : "()Ljava/lang/String;",
+                              null,
+                              null);
+                      if (!isAbstract) {
+                        method.visitCode();
+                        if (staticConflict) {
+                          method.visitInsn(ACONST_NULL);
+                        } else {
+                          method.visitLdcInsn("compatible");
+                        }
+                        method.visitInsn(ARETURN);
+                        method.visitMaxs(1, staticConflict ? 2 : 1);
+                      }
+                      method.visitEnd();
+                      super.visitEnd();
+                    }
+                  },
+                  0);
+          return writer.toByteArray();
+        });
+    FeatureFlaggingGateway.register(mock(FeatureFlaggingGateway.Backend.class));
+
+    try (URLClassLoader loader = sdkClassLoader()) {
+      final Class<?> contract = loader.loadClass(AgentBridge.class.getName());
+      final Object bridge = detect(loader);
+      if (isAbstract || staticConflict) {
+        assertSame(contract.getField("NONE").get(null), bridge);
+      } else {
+        assertEquals(JavaAgentBridge.class.getName(), bridge.getClass().getName());
+        assertEquals("compatible", contract.getMethod("additionalMethod").invoke(bridge));
+      }
+    }
+  }
+
+  private URLClassLoader sdkClassLoader() throws IOException {
+    return new URLClassLoader(
+        new URL[] {
+          sdkOverrides.toUri().toURL(),
+          AgentBridge.class.getProtectionDomain().getCodeSource().getLocation()
+        },
+        null);
+  }
+
+  private static Object detect(final ClassLoader loader) throws Exception {
+    return loader
+        .loadClass("com.datadog.openfeature.internal.bridge.AgentBridges")
+        .getMethod("detect")
+        .invoke(null);
+  }
+
+  private void rewriteSpi(final String spi, final UnaryOperator<byte[]> rewrite)
+      throws IOException {
+    final String resource = "com/datadog/openfeature/internal/bridge/" + spi + ".class";
+    final Path target = sdkOverrides.resolve(resource);
+    createDirectories(target.getParent());
+    try (InputStream input = AgentBridge.class.getClassLoader().getResourceAsStream(resource)) {
+      assertNotNull(input);
+      write(target, rewrite.apply(input.readAllBytes()));
+    }
+  }
+}

dd-java-agent/instrumentation/datadog/openfeature/dd-openfeature-1.69/src/test/java/datadog/trace/instrumentation/openfeature/DDOpenFeatureInstrumentationNoTracingForkedTest.java

@@ -0,0 +1,6 @@
+package datadog.trace.instrumentation.openfeature;
+
+import datadog.trace.test.junit.utils.config.WithConfig;
+
+@WithConfig(key = "trace.enabled", value = "false")
+class DDOpenFeatureInstrumentationNoTracingForkedTest extends DDOpenFeatureInstrumentationTest {}

dd-java-agent/instrumentation/datadog/openfeature/dd-openfeature-1.69/src/test/java/datadog/trace/instrumentation/openfeature/DDOpenFeatureInstrumentationTest.java

@@ -88,6 +88,17 @@ void postsThroughTheEventProxyOnlyWhenAvailable() throws Exception {
     assertThrows(EventProxyUnavailableException.class, () -> proxy.post("exposures", payload));
   }
 
+  @Test
+  void subscribesToTheAgentShutdown() throws Exception {
+    final FeatureFlaggingGateway.Backend backend = mock(FeatureFlaggingGateway.Backend.class);
+    final AutoCloseable subscription = mock(AutoCloseable.class);
+    final Runnable listener = mock(Runnable.class);
+    when(backend.subscribeShutdown(listener)).thenReturn(subscription);
+    FeatureFlaggingGateway.register(backend);
+
+    assertSame(subscription, AgentBridges.detect().shutdown().subscribe(listener));
+  }
+
   @Test
   void delegatesSpanEnrichment() {
     final FeatureFlaggingGateway.Backend backend = mock(FeatureFlaggingGateway.Backend.class);

dd-java-agent/instrumentation/datadog/openfeature/dd-openfeature-1.69/src/test/java/datadog/trace/instrumentation/openfeature/DDOpenFeatureIntegrationDisabledForkedTest.java

@@ -0,0 +1,26 @@
+package datadog.trace.instrumentation.openfeature;
+
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.mockito.Mockito.mock;
+
+import com.datadog.openfeature.internal.bridge.AgentBridge;
+import com.datadog.openfeature.internal.bridge.AgentBridges;
+import datadog.trace.agent.test.AbstractInstrumentationTest;
+import datadog.trace.api.featureflag.FeatureFlaggingGateway;
+import datadog.trace.test.junit.utils.config.WithConfig;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+
+@WithConfig(key = "trace.openfeature.enabled", value = "false")
+class DDOpenFeatureIntegrationDisabledForkedTest extends AbstractInstrumentationTest {
+  @AfterEach
+  void unregisterBackend() {
+    FeatureFlaggingGateway.register(null);
+  }
+
+  @Test
+  void doesNotInstrumentWhenTheIntegrationIsDisabled() {
+    FeatureFlaggingGateway.register(mock(FeatureFlaggingGateway.Backend.class));
+    assertSame(AgentBridge.NONE, AgentBridges.detect());
+  }
+}