Proposal 2: Deferred Commit Pattern in ASM Listener

dd-java-agent/appsec/src/main/java/com/datadog/appsec/config/AppSecConfigServiceImpl.java

@@ -64,6 +64,7 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -220,6 +221,11 @@ public String getCurrentRuleVersion() {
   }
 
   private class AppSecConfigChangesListener implements ProductListener {
+
+    // Deferred operations - collected during accept/remove, executed during commit
+    protected final Map<String, Map<String, Object>> configsToApply = new HashMap<>();
+    protected final Set<String> configsToRemove = new HashSet<>();
+
     @Override
     public void accept(ConfigKey configKey, byte[] content, PollingRateHinter pollingRateHinter)
         throws IOException {
@@ -235,13 +241,13 @@ public void accept(ConfigKey configKey, byte[] content, PollingRateHinter pollin
         ignoredConfigKeys.add(key);
       } else {
         ignoredConfigKeys.remove(key);
-        try {
-          beforeApply(key, contentMap);
-          maybeInitializeDefaultConfig();
-          handleWafUpdateResultReport(key, contentMap);
-        } catch (AppSecModule.AppSecModuleActivationException e) {
-          throw new RuntimeException(e);
-        }
+        maybeInitializeDefaultConfig();
+
+        // DEFER: Store config to apply in commit phase
+        configsToApply.put(key, contentMap);
+
+        // Cancel any pending remove for this key (accept overrides remove)
+        configsToRemove.remove(key);
       }
     }
 
@@ -252,43 +258,101 @@ public void remove(ConfigKey configKey, PollingRateHinter pollingRateHinter)
       if (ignoredConfigKeys.remove(key)) {
         return;
       }
-      try {
-        maybeInitializeDefaultConfig();
-        wafBuilder.removeConfig(key);
-        afterRemove(key);
-      } catch (UnclassifiedWafException e) {
-        throw new RuntimeException(e);
-      }
+
+      // DEFER: Mark config for removal in commit phase (don't execute immediately)
+      configsToRemove.add(key);
+
+      // Cancel any pending apply for this key (remove overrides apply)
+      configsToApply.remove(key);
     }
 
     @Override
     public void commit(PollingRateHinter pollingRateHinter) {
-      // no action needed
+      // Execute deferred operations atomically:
+      // 1. FIRST: Remove obsolete configs
+      // 2. THEN: Apply new/updated configs
+      // This ensures no duplicate rules in memory and no inconsistent state
+
+      // Phase 1: Execute all pending removes
+      for (String key : configsToRemove) {
+        if (shouldRemoveConfig(key)) {
+          try {
+            maybeInitializeDefaultConfig();
+            wafBuilder.removeConfig(key);
+          } catch (UnclassifiedWafException e) {
+            throw new RuntimeException(e);
+          }
+        }
+        afterRemove(key);
+      }
+
+      // Phase 2: Execute all pending applies
+      for (Map.Entry<String, Map<String, Object>> entry : configsToApply.entrySet()) {
+        beforeApply(entry.getKey(), entry.getValue());
+        try {
+          handleWafUpdateResultReport(entry.getKey(), entry.getValue());
+        } catch (AppSecModule.AppSecModuleActivationException e) {
+          throw new RuntimeException(e);
+        }
+      }
+
+      // Clear deferred operations after successful commit
+      configsToRemove.clear();
+      configsToApply.clear();
     }
 
     protected void beforeApply(final String key, final Map<String, Object> contentMap) {}
 
     protected void afterRemove(final String key) {}
+
+    /**
+     * Determines whether a config should be removed from the WAF. Override this method to add
+     * custom logic for specific product listeners.
+     *
+     * @param key the config key to check
+     * @return true if the config should be removed, false otherwise
+     */
+    protected boolean shouldRemoveConfig(final String key) {
+      return true; // Default: always remove
+    }
   }
 
   private class AppSecConfigChangesDDListener extends AppSecConfigChangesListener {
     @Override
     protected void beforeApply(final String key, final Map<String, Object> config) {
-      if (defaultConfigActivated) { // if we get any config, remove the default one
-        log.debug("Removing default config ASM_DD/default");
+      // Track that we're using this DD config key (for accounting)
+      usedDDWafConfigKeys.add(key);
+    }
+
+    @Override
+    protected void afterRemove(final String key) {
+      // Track removal from DD config keys (for accounting)
+      usedDDWafConfigKeys.remove(key);
+    }
+
+    @Override
+    protected boolean shouldRemoveConfig(final String key) {
+      // For ASM_DD, only remove if the config was actually added (tracked in usedDDWafConfigKeys)
+      return usedDDWafConfigKeys.contains(key);
+    }
+
+    @Override
+    public void commit(PollingRateHinter pollingRateHinter) {
+      // Special handling for ASM_DD: Remove default config before applying remote configs
+      // This must happen atomically with the other operations to avoid duplicate rules
+      if (defaultConfigActivated && !configsToApply.isEmpty()) {
+        log.debug("Removing default config ASM_DD/default before applying remote configs");
         try {
+          maybeInitializeDefaultConfig();
           wafBuilder.removeConfig(DEFAULT_WAF_CONFIG_RULE);
+          defaultConfigActivated = false;
         } catch (UnclassifiedWafException e) {
           throw new RuntimeException(e);
         }
-        defaultConfigActivated = false;
       }
-      usedDDWafConfigKeys.add(key);
-    }
 
-    @Override
-    protected void afterRemove(final String key) {
-      usedDDWafConfigKeys.remove(key);
+      // Execute the standard deferred commit flow (removes first, then applies)
+      super.commit(pollingRateHinter);
     }
   }
 

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/config/AppSecConfigServiceImplSpecification.groovy

@@ -731,12 +731,14 @@ class AppSecConfigServiceImplSpecification extends DDSpecification {
 
     when:
     listeners.savedWafDataChangesListener.accept(key, '''{"rules_override": [{"rules_target": [{"rule_id": "foo"}], "enabled": false}]}'''.getBytes(), NOOP)
+    listeners.savedWafDataChangesListener.commit(NOOP)
 
     then:
     service.usedDDWafConfigKeys.toList() == [key.toString()]
 
     when:
     listeners.savedWafDataChangesListener.remove(key, NOOP)
+    listeners.savedWafDataChangesListener.commit(NOOP)
 
     then:
     service.usedDDWafConfigKeys.empty
@@ -747,7 +749,6 @@ class AppSecConfigServiceImplSpecification extends DDSpecification {
     final key = new ParsedConfigKey('Test', '1234', 1, 'ASM_DD', 'ID')
 
     when:
-    AppSecSystem.active = true
     config.getAppSecActivation() >> ProductActivation.FULLY_ENABLED
     final service = new AppSecConfigServiceImpl(config, poller, reconf)
     service.init()
@@ -776,6 +777,118 @@ class AppSecConfigServiceImplSpecification extends DDSpecification {
     noExceptionThrown()
   }
 
+  void 'deferred commit: operations are tracked and executed via commit'() {
+    setup:
+    AppSecSystem.active = true
+    final key = new ParsedConfigKey('Test', '1234', 1, 'ASM_DD', 'ID')
+    final service = new AppSecConfigServiceImpl(config, poller, reconf)
+    config.getAppSecActivation() >> ProductActivation.ENABLED_INACTIVE
+
+    when:
+    service.maybeSubscribeConfigPolling()
+
+    then:
+    1 * poller.addListener(Product.ASM_DD, _) >> {
+      listeners.savedWafDataChangesListener = it[1]
+    }
+    1 * poller.addListener(Product.ASM_FEATURES, _, _) >> {
+      listeners.savedFeaturesDeserializer = it[1]
+      listeners.savedFeaturesListener = it[2]
+    }
+
+    when: 'activate AppSec'
+    listeners.savedFeaturesListener.accept('asm_features conf',
+      listeners.savedFeaturesDeserializer.deserialize('{"asm":{"enabled": true}}'.bytes),
+      NOOP)
+
+    then:
+    service.usedDDWafConfigKeys.empty
+
+    when: 'accept and commit'
+    listeners.savedWafDataChangesListener.accept(key, '''{"rules_override": [{"rules_target": [{"rule_id": "foo"}], "enabled": false}]}'''.getBytes(), NOOP)
+    listeners.savedWafDataChangesListener.commit(NOOP)
+
+    then:
+    service.usedDDWafConfigKeys.toList() == [key.toString()]
+
+    when: 'remove and commit'
+    listeners.savedWafDataChangesListener.remove(key, NOOP)
+    listeners.savedWafDataChangesListener.commit(NOOP)
+
+    then:
+    service.usedDDWafConfigKeys.empty
+  }
+
+  void 'deferred commit: accept after remove cancels the remove'() {
+    setup:
+    AppSecSystem.active = true
+    final key = new ParsedConfigKey('Test', '1234', 1, 'ASM_DD', 'ID')
+    final service = new AppSecConfigServiceImpl(config, poller, reconf)
+    config.getAppSecActivation() >> ProductActivation.ENABLED_INACTIVE
+
+    when:
+    service.maybeSubscribeConfigPolling()
+
+    then:
+    1 * poller.addListener(Product.ASM_DD, _) >> {
+      listeners.savedWafDataChangesListener = it[1]
+    }
+    1 * poller.addListener(Product.ASM_FEATURES, _, _) >> {
+      listeners.savedFeaturesDeserializer = it[1]
+      listeners.savedFeaturesListener = it[2]
+    }
+
+    when: 'activate AppSec and add config'
+    listeners.savedFeaturesListener.accept('asm_features conf',
+      listeners.savedFeaturesDeserializer.deserialize('{"asm":{"enabled": true}}'.bytes),
+      NOOP)
+    listeners.savedWafDataChangesListener.accept(key, '''{"rules_override": [{"rules_target": [{"rule_id": "foo"}], "enabled": false}]}'''.getBytes(), NOOP)
+    listeners.savedWafDataChangesListener.commit(NOOP)
+
+    then:
+    service.usedDDWafConfigKeys.toList() == [key.toString()]
+
+    when: 'remove then accept (cancel remove), then commit'
+    listeners.savedWafDataChangesListener.remove(key, NOOP)
+    listeners.savedWafDataChangesListener.accept(key, '''{"rules_override": [{"rules_target": [{"rule_id": "bar"}], "enabled": false}]}'''.getBytes(), NOOP)
+    listeners.savedWafDataChangesListener.commit(NOOP)
+
+    then: 'config updated not removed'
+    service.usedDDWafConfigKeys.toList() == [key.toString()]
+  }
+
+  void 'deferred commit: remove after accept cancels the accept'() {
+    setup:
+    AppSecSystem.active = true
+    final key = new ParsedConfigKey('Test', '1234', 1, 'ASM_DD', 'ID')
+    final service = new AppSecConfigServiceImpl(config, poller, reconf)
+    config.getAppSecActivation() >> ProductActivation.ENABLED_INACTIVE
+
+    when:
+    service.maybeSubscribeConfigPolling()
+
+    then:
+    1 * poller.addListener(Product.ASM_DD, _) >> {
+      listeners.savedWafDataChangesListener = it[1]
+    }
+    1 * poller.addListener(Product.ASM_FEATURES, _, _) >> {
+      listeners.savedFeaturesDeserializer = it[1]
+      listeners.savedFeaturesListener = it[2]
+    }
+
+    when: 'activate AppSec'
+    listeners.savedFeaturesListener.accept('asm_features conf',
+      listeners.savedFeaturesDeserializer.deserialize('{"asm":{"enabled": true}}'.bytes),
+      NOOP)
+
+    and: 'accept then remove (cancel accept), then commit'
+    listeners.savedWafDataChangesListener.accept(key, '''{"rules_override": [{"rules_target": [{"rule_id": "foo"}], "enabled": false}]}'''.getBytes(), NOOP)
+    listeners.savedWafDataChangesListener.remove(key, NOOP)
+    listeners.savedWafDataChangesListener.commit(NOOP)
+
+    then: 'config removed (never applied)'
+    service.usedDDWafConfigKeys.empty
+  }
 
   private static AppSecFeatures autoUserInstrum(String mode) {
     return new AppSecFeatures().tap { features ->

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/ddwaf/WAFModuleSpecification.groovy

@@ -145,10 +145,12 @@ class WAFModuleSpecification extends DDSpecification {
     ConfigKey config = new ParsedConfigKey(configKey, 'null', 1, 'null', 'null')
     if(map == null) {
       listener.remove(config, null)
-      return
+    } else {
+      def json = ADAPTER.toJson(map)
+      listener.accept(config, json.getBytes(), null)
     }
-    def json = ADAPTER.toJson(map)
-    listener.accept(config, json.getBytes(), null)
+    // Trigger commit to execute deferred operations
+    listener.commit(null)
   }
 
 

remote-config/remote-config-core/src/main/java/datadog/remoteconfig/state/ProductState.java

@@ -58,10 +58,8 @@ public boolean apply(
     errors = null;
 
     List<ParsedConfigKey> configBeenUsedByProduct = new ArrayList<>();
-    List<ParsedConfigKey> changedKeys = new ArrayList<>();
     boolean changesDetected = false;
 
-    // Step 1: Detect all changes
     for (ParsedConfigKey configKey : relevantKeys) {
       try {
         RemoteConfigResponse.Targets.ConfigTarget target =
@@ -70,28 +68,14 @@ public boolean apply(
 
         if (isTargetChanged(configKey, target)) {
           changesDetected = true;
-          changedKeys.add(configKey);
-        }
-      } catch (ReportableException e) {
-        recordError(e);
-      }
-    }
-
-    // Step 2: For products other than ASM_DD, apply changes immediately
-    if (product != Product.ASM_DD) {
-      for (ParsedConfigKey configKey : changedKeys) {
-        try {
           byte[] content = getTargetFileContent(fleetResponse, configKey);
           callListenerApplyTarget(fleetResponse, hinter, configKey, content);
-        } catch (ReportableException e) {
-          recordError(e);
         }
+      } catch (ReportableException e) {
+        recordError(e);
       }
     }
 
-    // Step 3: Remove obsolete configurations (for all products)
-    // For ASM_DD, this is critical: removes MUST happen before applies to prevent
-    // duplicate rule warnings from the ddwaf rule parser and causing memory spikes.
     List<ParsedConfigKey> keysToRemove =
         cachedTargetFiles.keySet().stream()
             .filter(configKey -> !configBeenUsedByProduct.contains(configKey))
@@ -102,22 +86,6 @@ public boolean apply(
       callListenerRemoveTarget(hinter, configKey);
     }
 
-    // Step 4: For ASM_DD, apply changes AFTER removes
-    // TODO: This is a temporary solution. The proper fix requires better synchronization
-    // between remove and add/update operations. This should be discussed
-    // with the guild to determine the best long-term design approach.
-    if (product == Product.ASM_DD) {
-      for (ParsedConfigKey configKey : changedKeys) {
-        try {
-          byte[] content = getTargetFileContent(fleetResponse, configKey);
-          callListenerApplyTarget(fleetResponse, hinter, configKey, content);
-        } catch (ReportableException e) {
-          recordError(e);
-        }
-      }
-    }
-
-    // Step 5: Commit if there were changes
     if (changesDetected) {
       try {
         callListenerCommit(hinter);