diff --git a/modules/codegen/src/main/java/org/apache/ignite/internal/MessageSerializerGenerator.java b/modules/codegen/src/main/java/org/apache/ignite/internal/MessageSerializerGenerator.java
index 6ae7e0294109d..81680f1b1088c 100644
--- a/modules/codegen/src/main/java/org/apache/ignite/internal/MessageSerializerGenerator.java
+++ b/modules/codegen/src/main/java/org/apache/ignite/internal/MessageSerializerGenerator.java
@@ -23,6 +23,7 @@
import java.util.ArrayList;
import java.util.BitSet;
import java.util.Collection;
+import java.util.EnumSet;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
@@ -86,6 +87,9 @@ public class MessageSerializerGenerator extends MessageCompanionGenerator {
/** */
private static final String MESSAGE_COLLECTION_TYPE_CLS = "org.apache.ignite.plugin.extensions.communication.MessageCollectionType";
+ /** */
+ private static final String COLLECTION_IMPL_TYPE_CLS = "org.apache.ignite.plugin.extensions.communication.CollectionImplementationType";
+
/** */
private static final String IGNITE_UUID_CLS = "org.apache.ignite.lang.IgniteUuid";
@@ -552,9 +556,17 @@ else if (assignableFrom(erasedType(type), type(Collection.class.getName()))) {
assert typeArgs.size() == 1 : type.toString();
- return "new MessageCollectionType(" +
- messageCollectionItemTypeDescriptor(typeArgs.get(0), field) + ", " +
- assignableFrom(erasedType(type), type(Set.class.getName())) + ")";
+ TypeMirror itemType = typeArgs.get(0);
+
+ if (EnumSet.class.getName().equals(qualifiedClassName(erasedType(type))) && !enumType(env, itemType))
+ throw new IllegalArgumentException("Unexpected Enum Set element type [itemType=" + itemType + ", colType=" + type + ']');
+
+ imports.add(COLLECTION_IMPL_TYPE_CLS);
+
+ String implType = resolveCollectionImplementationType(type);
+
+ return "new MessageCollectionType(" + messageCollectionItemTypeDescriptor(itemType, field) +
+ ", CollectionImplementationType." + implType + ")";
}
else if (enumType(env, type)) {
imports.add("org.apache.ignite.plugin.extensions.communication.MessageEnumType");
@@ -566,7 +578,7 @@ else if (enumType(env, type)) {
String decoder = custMapper ? prefix + "Mapper::decode" :
"b -> DefaultEnumMapper.INSTANCE.decode(" + prefix + "Vals, b)";
- return String.format("new MessageEnumType<>(%s, %s)", encoder, decoder);
+ return String.format("new MessageEnumType<>(%s.class, %s, %s)", simpleClassName(type), encoder, decoder);
}
else {
imports.add(MESSAGE_ITEM_TYPE_CLS);
@@ -575,6 +587,21 @@ else if (enumType(env, type)) {
}
}
+ /**
+ * @param type Declared collection type.
+ * @return Name of the {@code CollectionImplementationType} that will be used to create collection instance.
+ */
+ private String resolveCollectionImplementationType(TypeMirror type) {
+ TypeMirror declType = erasedType(type);
+
+ if (EnumSet.class.getName().equals(qualifiedClassName(declType)))
+ return "ENUM_SET";
+ else if (assignableFrom(declType, type(Set.class.getName())))
+ return "HASH_SET";
+ else
+ return "ARRAY_LIST";
+ }
+
/**
* Find MessageCollectionItemType for a given type.
*
diff --git a/modules/core/src/main/java/org/apache/ignite/internal/direct/stream/DirectByteBufferStream.java b/modules/core/src/main/java/org/apache/ignite/internal/direct/stream/DirectByteBufferStream.java
index 75440117e23db..c0752343b5441 100644
--- a/modules/core/src/main/java/org/apache/ignite/internal/direct/stream/DirectByteBufferStream.java
+++ b/modules/core/src/main/java/org/apache/ignite/internal/direct/stream/DirectByteBufferStream.java
@@ -22,6 +22,7 @@
import java.util.ArrayList;
import java.util.BitSet;
import java.util.Collection;
+import java.util.EnumSet;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
@@ -1639,11 +1640,11 @@ public T[] readObjectArray(MessageArrayType type, MessageReader reader) {
}
/**
- * Reads collection eather as a {@link ArrayList} or a {@link HashSet}.
+ * Reads collection either as an {@link ArrayList}, a {@link HashSet} or an {@link EnumSet}.
*
* @param type Item type.
* @param reader Reader.
- * @return {@link ArrayList} or a {@link HashSet}.
+ * @return {@link ArrayList}, {@link HashSet} or {@link EnumSet}.
*/
public > C readCollection(MessageCollectionType type, MessageReader reader) {
if (readSize == -1) {
@@ -1657,7 +1658,7 @@ public > C readCollection(MessageCollectionType type, Me
if (readSize >= 0) {
if (col == null)
- col = type.set() ? U.newHashSet(readSize) : new ArrayList<>(readSize);
+ col = newCollection(type);
for (int i = readItems; i < readSize; i++) {
Object item = read(type.valueType(), reader);
@@ -1682,6 +1683,16 @@ public > C readCollection(MessageCollectionType type, Me
return col0;
}
+ /** */
+ @SuppressWarnings("unchecked")
+ private Collection