IGNITE-28520 Move prepareMarshal / finishUnmarshal out of NIO communication thread

Author: abashev

modules/codegen/src/main/java/org/apache/ignite/internal/MessageSerializerGenerator.java

@@ -114,6 +114,12 @@ public class MessageSerializerGenerator {
     /** The marshallable message type. */
     private final TypeMirror marshallableMsgType;
 
+    /** Collection of lines for the {@code prepareMarshalCacheObjects} method. Empty when the message has no such fields. */
+    private final List<String> prepareCacheObjects = new ArrayList<>();
+
+    /** Collection of lines for the {@code finishUnmarshalCacheObjects} method. Empty when the message has no such fields. */
+    private final List<String> finishCacheObjects = new ArrayList<>();
+
     /** */
     private int indent;
 
@@ -189,6 +195,20 @@ private String generateSerializerCode(String serClsName) throws IOException {
 
             writer.write(TAB + "}" + NL);
 
+            // Write #prepareMarshalCacheObjects / #finishUnmarshalCacheObjects overrides if the message has any
+            // cache-object-bearing @Order fields (IGNITE-28520).
+            if (!prepareCacheObjects.isEmpty()) {
+                writer.write(NL);
+
+                for (String p: prepareCacheObjects)
+                    writer.write(p + NL);
+
+                writer.write(NL);
+
+                for (String f: finishCacheObjects)
+                    writer.write(f + NL);
+            }
+
             writer.write("}");
 
             writer.write(NL);
@@ -242,6 +262,178 @@ private void generateMethods(List<VariableElement> fields) throws Exception {
 
         finish(write, false, false);
         finish(read, true, marshallableMessage());
+
+        generateCacheObjectMethods(fields);
+    }
+
+    /**
+     * Generates bodies of {@code prepareMarshalCacheObjects} / {@code finishUnmarshalCacheObjects} overrides when the
+     * message carries at least one {@code @Order}-annotated field whose declared type is {@code CacheObject},
+     * {@code KeyCacheObject}, a {@code Collection} of them, or an array of them. Traversal is single-level: it does
+     * not recurse into nested messages or enter map values. See IGNITE-28520.
+     */
+    private void generateCacheObjectMethods(List<VariableElement> fields) throws Exception {
+        List<VariableElement> cacheObjFields = new ArrayList<>();
+
+        for (VariableElement field: fields) {
+            if (isCacheObjectBearingField(field.asType()))
+                cacheObjFields.add(field);
+        }
+
+        if (cacheObjFields.isEmpty())
+            return;
+
+        imports.add("org.apache.ignite.IgniteCheckedException");
+        imports.add("org.apache.ignite.internal.processors.cache.CacheObjectValueContext");
+
+        startCacheObjectMethod(prepareCacheObjects, true);
+        startCacheObjectMethod(finishCacheObjects, false);
+
+        indent++;
+
+        for (VariableElement field: cacheObjFields) {
+            emitCacheObjectCall(prepareCacheObjects, field, true);
+            emitCacheObjectCall(finishCacheObjects, field, false);
+        }
+
+        indent--;
+
+        prepareCacheObjects.add(identedLine("}"));
+        finishCacheObjects.add(identedLine("}"));
+    }
+
+    /**
+     * @return {@code true} if the given field type is a {@code CacheObject}/{@code KeyCacheObject}, a collection of
+     *      them, or an array of them. Maps and nested messages are intentionally excluded — see IGNITE-28520 scope.
+     */
+    private boolean isCacheObjectBearingField(TypeMirror type) {
+        if (type.getKind() == TypeKind.ARRAY) {
+            TypeMirror comp = ((ArrayType)type).getComponentType();
+
+            return comp.getKind() == TypeKind.DECLARED && isCacheObjectType(comp);
+        }
+
+        if (type.getKind() != TypeKind.DECLARED)
+            return false;
+
+        if (isCacheObjectType(type))
+            return true;
+
+        // Map is intentionally skipped even though Map is Collection-unrelated — exclude it explicitly for clarity.
+        if (assignableFrom(erasedType(type), type(Map.class.getName())))
+            return false;
+
+        if (assignableFrom(erasedType(type), type(Collection.class.getName()))) {
+            List<? extends TypeMirror> typeArgs = ((DeclaredType)type).getTypeArguments();
+
+            if (typeArgs.size() == 1) {
+                TypeMirror arg = typeArgs.get(0);
+
+                return arg.getKind() == TypeKind.DECLARED && isCacheObjectType(arg);
+            }
+        }
+
+        return false;
+    }
+
+    /** @return {@code true} if {@code type} is assignable to {@code CacheObject} (this also covers {@code KeyCacheObject}). */
+    private boolean isCacheObjectType(TypeMirror type) {
+        return assignableFrom(type, type("org.apache.ignite.internal.processors.cache.CacheObject"));
+    }
+
+    /** Emits method signature and opening brace for cache-object marshalling methods. */
+    private void startCacheObjectMethod(List<String> code, boolean prepare) {
+        indent = 1;
+
+        code.add(identedLine(METHOD_JAVADOC));
+
+        if (prepare) {
+            code.add(identedLine(
+                "@Override public void prepareMarshalCacheObjects(" + type.getSimpleName() +
+                    " msg, CacheObjectValueContext ctx) throws IgniteCheckedException {"));
+        }
+        else {
+            code.add(identedLine(
+                "@Override public void finishUnmarshalCacheObjects(" + type.getSimpleName() +
+                    " msg, CacheObjectValueContext ctx, ClassLoader ldr) throws IgniteCheckedException {"));
+        }
+    }
+
+    /**
+     * Emits a single field traversal statement — a guarded {@code prepareMarshal} / {@code finishUnmarshal} call
+     * for a direct {@code CacheObject} field, or a null-safe loop for a collection / array of them.
+     */
+    private void emitCacheObjectCall(List<String> code, VariableElement field, boolean prepare) {
+        String mtd = prepare ? "prepareMarshal(ctx)" : "finishUnmarshal(ctx, ldr)";
+
+        String accessor = fieldAccessor(field);
+
+        TypeMirror type = field.asType();
+
+        if (type.getKind() == TypeKind.DECLARED && isCacheObjectType(type)) {
+            code.add(identedLine("if (%s != null)", accessor));
+
+            indent++;
+
+            code.add(identedLine("%s.%s;", accessor, mtd));
+
+            indent--;
+        }
+        else {
+            // Collection or array of CacheObject — iterate with a null-check on both the container and the element.
+            code.add(identedLine("if (%s != null) {", accessor));
+
+            indent++;
+
+            String elementType = "org.apache.ignite.internal.processors.cache.KeyCacheObject";
+
+            if (type.getKind() == TypeKind.ARRAY) {
+                TypeMirror comp = ((ArrayType)type).getComponentType();
+
+                if (!assignableFrom(comp, type(elementType)))
+                    elementType = "org.apache.ignite.internal.processors.cache.CacheObject";
+            }
+            else {
+                TypeMirror arg = ((DeclaredType)type).getTypeArguments().get(0);
+
+                if (!assignableFrom(arg, type(elementType)))
+                    elementType = "org.apache.ignite.internal.processors.cache.CacheObject";
+            }
+
+            imports.add(elementType);
+
+            String simpleElement = elementType.substring(elementType.lastIndexOf('.') + 1);
+
+            code.add(identedLine("for (%s obj : %s) {", simpleElement, accessor));
+
+            indent++;
+
+            code.add(identedLine("if (obj != null)"));
+
+            indent++;
+
+            code.add(identedLine("obj.%s;", mtd));
+
+            indent--;
+
+            indent--;
+
+            code.add(identedLine("}"));
+
+            indent--;
+
+            code.add(identedLine("}"));
+        }
+    }
+
+    /** Returns the field access expression, taking superclass-field access into account. */
+    private String fieldAccessor(VariableElement field) {
+        String name = field.getSimpleName().toString();
+
+        if (type.equals(field.getEnclosingElement()))
+            return "msg." + name;
+
+        return "((" + field.getEnclosingElement().getSimpleName() + ")msg)." + name;
     }
 
     /**

modules/core/src/main/java/org/apache/ignite/internal/processors/cache/GridCacheIoManager.java

@@ -41,6 +41,7 @@
 import org.apache.ignite.internal.cluster.ClusterTopologyCheckedException;
 import org.apache.ignite.internal.managers.communication.GridIoPolicy;
 import org.apache.ignite.internal.managers.communication.GridMessageListener;
+import org.apache.ignite.plugin.extensions.communication.MessageSerializer;
 import org.apache.ignite.internal.managers.deployment.GridDeploymentInfo;
 import org.apache.ignite.internal.processors.affinity.AffinityTopologyVersion;
 import org.apache.ignite.internal.processors.cache.distributed.dht.CacheGetFuture;
@@ -1150,13 +1151,81 @@ private boolean onSend(GridCacheMessage msg, @Nullable UUID destNodeId) throws I
         if (destNodeId == null || !cctx.localNodeId().equals(destNodeId)) {
             msg.prepareMarshal(cctx);
 
+            prepareMarshalGeneratedCacheObjects(msg);
+
             if (msg instanceof GridCacheDeployable && msg.addDeploymentInfo())
                 cctx.deploy().prepare((GridCacheDeployable)msg);
         }
 
         return true;
     }
 
+    /**
+     * Invokes {@link MessageSerializer#prepareMarshalCacheObjects} on the user thread for the given message, as part
+     * of the IGNITE-28520 Phase 1 two-phase marshalling. The generated serializer method walks {@code @Order}-annotated
+     * {@link CacheObject}/{@link org.apache.ignite.internal.processors.cache.KeyCacheObject} fields of {@code msg}
+     * (including single-level collections and arrays) and calls
+     * {@link CacheObject#prepareMarshal(CacheObjectValueContext)} on each, so that the NIO communication thread never
+     * has to do this work.
+     *
+     * <p>This call is intentionally <b>additive</b> to the existing
+     * {@link GridCacheMessage#prepareMarshal(GridCacheSharedContext)} override path: it runs right after the manual
+     * override. For messages whose override has already populated {@code valBytes} of every cache object field, this
+     * call becomes a no-op thanks to the idempotency guard inside
+     * {@link org.apache.ignite.internal.processors.cache.CacheObjectImpl#prepareMarshal(CacheObjectValueContext)}
+     * ({@code if (valBytes == null)}). For messages that don't override {@code prepareMarshal(GridCacheSharedContext)}
+     * but do carry cache-object fields, the generated method is what actually does the marshalling work.
+     *
+     * <p>Only single-cache messages ({@link GridCacheIdMessage}) are processed here — cross-cache messages resolve
+     * their per-entry cache contexts internally via their hand-written overrides, and we don't try to second-guess
+     * them from a single shared context.
+     *
+     * @param msg Message being sent.
+     * @throws IgniteCheckedException If the generated serializer fails to marshal a cache object field.
+     */
+    @SuppressWarnings({"unchecked", "rawtypes"})
+    private void prepareMarshalGeneratedCacheObjects(GridCacheMessage msg) throws IgniteCheckedException {
+        if (!(msg instanceof GridCacheIdMessage))
+            return;
+
+        CacheObjectContext cacheObjCtx = cctx.cacheObjectContext(((GridCacheIdMessage)msg).cacheId());
+
+        if (cacheObjCtx == null)
+            return;
+
+        MessageSerializer ser = cctx.gridIO().messageFactory().serializer(msg.directType());
+
+        if (ser != null)
+            ser.prepareMarshalCacheObjects(msg, cacheObjCtx);
+    }
+
+    /**
+     * Invokes {@link MessageSerializer#finishUnmarshalCacheObjects} on the user (listener-dispatch) thread — the
+     * receive-side mirror of {@link #prepareMarshalGeneratedCacheObjects}. See that method's docs for the full
+     * contract. This call is also additive to the manual
+     * {@link GridCacheMessage#finishUnmarshal(GridCacheSharedContext, ClassLoader)} override path and idempotent for
+     * cache objects whose values have already been deserialized.
+     *
+     * @param msg Message being received.
+     * @param ldr Class loader used for value deserialization.
+     * @throws IgniteCheckedException If the generated serializer fails to unmarshal a cache object field.
+     */
+    @SuppressWarnings({"unchecked", "rawtypes"})
+    private void finishUnmarshalGeneratedCacheObjects(GridCacheMessage msg, ClassLoader ldr) throws IgniteCheckedException {
+        if (!(msg instanceof GridCacheIdMessage))
+            return;
+
+        CacheObjectContext cacheObjCtx = cctx.cacheObjectContext(((GridCacheIdMessage)msg).cacheId());
+
+        if (cacheObjCtx == null)
+            return;
+
+        MessageSerializer ser = cctx.gridIO().messageFactory().serializer(msg.directType());
+
+        if (ser != null)
+            ser.finishUnmarshalCacheObjects(msg, cacheObjCtx, ldr);
+    }
+
     /**
      * @param nodeId Node ID.
      * @param sndErr Send error.
@@ -1604,7 +1673,11 @@ private void unmarshall(UUID nodeId, GridCacheMessage cacheMsg) {
                     log.debug("Set P2P context [senderId=" + nodeId + ", msg=" + cacheMsg + ']');
             }
 
-            cacheMsg.finishUnmarshal(cctx, cctx.deploy().globalLoader());
+            ClassLoader ldr = cctx.deploy().globalLoader();
+
+            cacheMsg.finishUnmarshal(cctx, ldr);
+
+            finishUnmarshalGeneratedCacheObjects(cacheMsg, ldr);
         }
         catch (IgniteCheckedException e) {
             cacheMsg.onClassError(e);

modules/core/src/main/java/org/apache/ignite/plugin/extensions/communication/MessageSerializer.java

@@ -17,6 +17,9 @@
 
 package org.apache.ignite.plugin.extensions.communication;
 
+import org.apache.ignite.IgniteCheckedException;
+import org.apache.ignite.internal.processors.cache.CacheObjectValueContext;
+
 /**
  * Interface for message serialization logic.
  */
@@ -38,4 +41,44 @@ public interface MessageSerializer<M extends Message> {
      * @return Whether message was fully read.
      */
     public boolean readFrom(M msg, MessageReader reader);
+
+    /**
+     * Prepares {@link org.apache.ignite.internal.processors.cache.CacheObject} fields of the given message for
+     * marshalling. This is the first phase of the two-phase marshalling introduced by IGNITE-28520: it is expected
+     * to be invoked on a <b>user thread</b> before the message is handed off to the NIO worker, so that
+     * {@link org.apache.ignite.internal.processors.cache.CacheObject#prepareMarshal(CacheObjectValueContext)} never
+     * runs on the NIO communication thread (or, worse, the Discovery thread where it could deadlock).
+     *
+     * <p>The default implementation is a no-op. The generated serializer overrides it for messages that contain
+     * at least one {@code @Order}-annotated field whose declared type is {@code CacheObject} / {@code KeyCacheObject}
+     * (including {@code Collection<CacheObject>} and {@code CacheObject[]}). Only the top level is traversed — the
+     * traversal does <b>not</b> recurse into nested messages or map values.
+     *
+     * @param msg Message instance.
+     * @param ctx Cache object value context.
+     * @throws IgniteCheckedException If marshalling of any contained cache object fails.
+     */
+    public default void prepareMarshalCacheObjects(M msg, CacheObjectValueContext ctx) throws IgniteCheckedException {
+        // No-op by default.
+    }
+
+    /**
+     * Finishes unmarshalling of {@link org.apache.ignite.internal.processors.cache.CacheObject} fields of the given
+     * message. Intended to be invoked on a <b>user thread</b> (for example, on a listener dispatch thread) after
+     * {@link #readFrom(Message, MessageReader)} returns {@code true} on the NIO worker, so that
+     * {@link org.apache.ignite.internal.processors.cache.CacheObject#finishUnmarshal(CacheObjectValueContext, ClassLoader)}
+     * never runs on the NIO communication thread.
+     *
+     * <p>The default implementation is a no-op. Mirrors the traversal rules of
+     * {@link #prepareMarshalCacheObjects(Message, CacheObjectValueContext)}.
+     *
+     * @param msg Message instance.
+     * @param ctx Cache object value context.
+     * @param ldr Class loader to use for value deserialization.
+     * @throws IgniteCheckedException If unmarshalling of any contained cache object fails.
+     */
+    public default void finishUnmarshalCacheObjects(M msg, CacheObjectValueContext ctx, ClassLoader ldr)
+        throws IgniteCheckedException {
+        // No-op by default.
+    }
 }