WFLY-9146 Fix transformation of legacy protocol write-attribute opera…

…tions

clustering/common/src/main/java/org/jboss/as/clustering/controller/OperationStepHandlerDescriptor.java

@@ -25,9 +25,7 @@
 import java.util.Collections;
 import java.util.Map;
 import java.util.function.Predicate;
-import java.util.function.UnaryOperator;
 
-import org.jboss.as.controller.OperationStepHandler;
 import org.jboss.dmr.ModelNode;
 
 /**
@@ -42,13 +40,4 @@ public interface OperationStepHandlerDescriptor {
     default Map<Capability, Predicate<ModelNode>> getCapabilities() {
         return Collections.emptyMap();
     }
-
-    /**
-     * Returns a transformer to be applied to all operations that operate on an existing resource.
-     * This is typically used to adapt legacy operations to conform to the current version of the model.
-     * @return an operation handler transformer.
-     */
-    default UnaryOperator<OperationStepHandler> getOperationTransformation() {
-        return UnaryOperator.identity();
-    }
 }

clustering/common/src/main/java/org/jboss/as/clustering/controller/RemoveStepHandlerDescriptor.java

@@ -25,8 +25,10 @@
 import java.util.Collections;
 import java.util.Map;
 import java.util.function.Function;
+import java.util.function.UnaryOperator;
 
 import org.jboss.as.controller.CapabilityReferenceRecorder;
+import org.jboss.as.controller.OperationStepHandler;
 import org.jboss.as.controller.PathAddress;
 import org.jboss.as.controller.descriptions.ResourceDescriptionResolver;
 
@@ -58,4 +60,13 @@ default Collection<RuntimeResourceRegistration> getRuntimeResourceRegistrations(
     default Map<CapabilityReferenceRecorder, Function<PathAddress, String>> getResourceCapabilityReferences() {
         return Collections.emptyMap();
     }
+
+    /**
+     * Returns a transformer to be applied to all operations that operate on an existing resource.
+     * This is typically used to adapt legacy operations to conform to the current version of the model.
+     * @return an operation handler transformer.
+     */
+    default UnaryOperator<OperationStepHandler> getOperationTransformation() {
+        return UnaryOperator.identity();
+    }
 }

clustering/common/src/main/java/org/jboss/as/clustering/controller/ResourceRegistration.java

@@ -22,11 +22,14 @@
 
 package org.jboss.as.clustering.controller;
 
-import java.util.function.UnaryOperator;
+import java.util.Optional;
 
+import org.jboss.as.clustering.function.Predicates;
+import org.jboss.as.controller.OperationDefinition;
 import org.jboss.as.controller.OperationStepHandler;
 import org.jboss.as.controller.operations.global.ListOperations;
 import org.jboss.as.controller.operations.global.MapOperations;
+import org.jboss.as.controller.operations.global.WriteAttributeHandler;
 import org.jboss.as.controller.registry.ManagementResourceRegistration;
 
 /**
@@ -64,16 +67,24 @@ public void register(ManagementResourceRegistration registration) {
         this.addRegistration.register(registration);
         this.removeRegistration.register(registration);
 
-        UnaryOperator<OperationStepHandler> transformer = this.descriptor.getOperationTransformation();
+        // Override global operations with transformed operations, if necessary
+        this.registerTransformedOperation(registration, WriteAttributeHandler.DEFINITION, WriteAttributeHandler.INSTANCE);
 
-        registration.registerOperationHandler(MapOperations.MAP_PUT_DEFINITION, transformer.apply(MapOperations.MAP_PUT_HANDLER));
-        registration.registerOperationHandler(MapOperations.MAP_GET_DEFINITION, transformer.apply(MapOperations.MAP_GET_HANDLER));
-        registration.registerOperationHandler(MapOperations.MAP_REMOVE_DEFINITION, transformer.apply(MapOperations.MAP_REMOVE_HANDLER));
-        registration.registerOperationHandler(MapOperations.MAP_CLEAR_DEFINITION, transformer.apply(MapOperations.MAP_CLEAR_HANDLER));
+        this.registerTransformedOperation(registration, MapOperations.MAP_PUT_DEFINITION, MapOperations.MAP_PUT_HANDLER);
+        this.registerTransformedOperation(registration, MapOperations.MAP_GET_DEFINITION, MapOperations.MAP_GET_HANDLER);
+        this.registerTransformedOperation(registration, MapOperations.MAP_REMOVE_DEFINITION, MapOperations.MAP_REMOVE_HANDLER);
+        this.registerTransformedOperation(registration, MapOperations.MAP_CLEAR_DEFINITION, MapOperations.MAP_CLEAR_HANDLER);
 
-        registration.registerOperationHandler(ListOperations.LIST_ADD_DEFINITION, transformer.apply(ListOperations.LIST_ADD_HANDLER));
-        registration.registerOperationHandler(ListOperations.LIST_REMOVE_DEFINITION, transformer.apply(ListOperations.LIST_REMOVE_HANDLER));
-        registration.registerOperationHandler(ListOperations.LIST_GET_DEFINITION, transformer.apply(ListOperations.LIST_GET_HANDLER));
-        registration.registerOperationHandler(ListOperations.LIST_CLEAR_DEFINITION, transformer.apply(ListOperations.LIST_CLEAR_HANDLER));
+        this.registerTransformedOperation(registration, ListOperations.LIST_ADD_DEFINITION, ListOperations.LIST_ADD_HANDLER);
+        this.registerTransformedOperation(registration, ListOperations.LIST_GET_DEFINITION, ListOperations.LIST_GET_HANDLER);
+        this.registerTransformedOperation(registration, ListOperations.LIST_REMOVE_DEFINITION, ListOperations.LIST_REMOVE_HANDLER);
+        this.registerTransformedOperation(registration, ListOperations.LIST_CLEAR_DEFINITION, ListOperations.LIST_CLEAR_HANDLER);
+    }
+
+    private void registerTransformedOperation(ManagementResourceRegistration registration, OperationDefinition definition, OperationStepHandler handler) {
+        // Only override global operation handlers for non-identity transformations
+        Optional.of(handler).map(this.descriptor.getOperationTransformation())
+                .filter(Predicates.same(handler).negate())
+                .ifPresent(transformedHandler -> registration.registerOperationHandler(definition, transformedHandler));
     }
 }

clustering/common/src/main/java/org/jboss/as/clustering/controller/RestartParentResourceWriteAttributeHandler.java

@@ -22,12 +22,9 @@
 
 package org.jboss.as.clustering.controller;
 
-import java.util.function.UnaryOperator;
-
 import org.jboss.as.clustering.controller.transform.InitialAttributeValueOperationContextAttachment;
 import org.jboss.as.controller.OperationContext;
 import org.jboss.as.controller.OperationFailedException;
-import org.jboss.as.controller.OperationStepHandler;
 import org.jboss.as.controller.PathAddress;
 import org.jboss.as.controller.RestartParentWriteAttributeHandler;
 import org.jboss.as.controller.registry.ManagementResourceRegistration;
@@ -68,8 +65,7 @@ protected PathAddress getParentAddress(PathAddress address) {
 
     @Override
     public void register(ManagementResourceRegistration registration) {
-        UnaryOperator<OperationStepHandler> transformer = this.descriptor.getOperationTransformation();
-        this.descriptor.getAttributes().forEach(attribute -> registration.registerReadWriteAttribute(attribute, null, transformer.apply(this)));
+        this.descriptor.getAttributes().forEach(attribute -> registration.registerReadWriteAttribute(attribute, null, this));
     }
 
     @Override

clustering/common/src/main/java/org/jboss/as/clustering/controller/WriteAttributeStepHandler.java

@@ -24,13 +24,11 @@
 
 import java.util.Map;
 import java.util.function.Predicate;
-import java.util.function.UnaryOperator;
 
 import org.jboss.as.clustering.controller.transform.InitialAttributeValueOperationContextAttachment;
 import org.jboss.as.controller.AttributeDefinition;
 import org.jboss.as.controller.OperationContext;
 import org.jboss.as.controller.OperationFailedException;
-import org.jboss.as.controller.OperationStepHandler;
 import org.jboss.as.controller.PathAddress;
 import org.jboss.as.controller.ReloadRequiredWriteAttributeHandler;
 import org.jboss.as.controller.registry.AttributeAccess;
@@ -65,8 +63,7 @@ protected boolean requiresRuntime(OperationContext context) {
 
     @Override
     public void register(ManagementResourceRegistration registration) {
-        UnaryOperator<OperationStepHandler> transformer = this.descriptor.getOperationTransformation();
-        this.descriptor.getAttributes().forEach(attribute -> registration.registerReadWriteAttribute(attribute, null, transformer.apply(this)));
+        this.descriptor.getAttributes().forEach(attribute -> registration.registerReadWriteAttribute(attribute, null, this));
     }
 
     @Override

clustering/common/src/main/java/org/jboss/as/clustering/function/Predicates.java

@@ -41,4 +41,8 @@ public static <T> Predicate<T> never() {
     public static <T> Predicate<T> when(boolean condition) {
         return test -> condition;
     }
+
+    public static <T> Predicate<T> same(T object) {
+        return test -> test == object;
+    }
 }