Improve MessageKryoRegistrar for registrations

(cherry picked from commit 5ec71d4)

Author: artembilan

spring-integration-core/src/main/java/org/springframework/integration/codec/kryo/CompositeKryoRegistrar.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2015-2019 the original author or authors.
+ * Copyright 2015-2020 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -22,6 +22,7 @@
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
 
+import com.esotericsoftware.kryo.Kryo;
 import com.esotericsoftware.kryo.Registration;
 
 /**
@@ -42,6 +43,13 @@ public CompositeKryoRegistrar(List<KryoRegistrar> delegates) {
 		}
 	}
 
+	@Override
+	public void registerTypes(Kryo kryo) {
+		for (KryoRegistrar registrar : this.delegates) {
+			registrar.registerTypes(kryo);
+		}
+	}
+
 	@Override
 	public final List<Registration> getRegistrations() {
 		List<Registration> registrations = new ArrayList<Registration>();

spring-integration-core/src/main/java/org/springframework/integration/codec/kryo/MessageCodec.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2015-2019 the original author or authors.
+ * Copyright 2015-2020 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -18,7 +18,9 @@
 
 /**
  * {@link PojoCodec} configured to encode/decode {@code Message<?>}s.
+ *
  * @author Gary Russell
+ *
  * @since 4.2
  *
  */

spring-integration-core/src/main/java/org/springframework/integration/codec/kryo/MessageKryoRegistrar.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2015-2019 the original author or authors.
+ * Copyright 2015-2020 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,41 +17,127 @@
 package org.springframework.integration.codec.kryo;
 
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.List;
+import java.util.UUID;
 
+import org.springframework.integration.message.AdviceMessage;
+import org.springframework.integration.support.MutableMessage;
 import org.springframework.integration.support.MutableMessageHeaders;
 import org.springframework.messaging.MessageHeaders;
+import org.springframework.messaging.support.ErrorMessage;
+import org.springframework.messaging.support.GenericMessage;
 
+import com.esotericsoftware.kryo.Kryo;
 import com.esotericsoftware.kryo.Registration;
 
 /**
  * Registers common MessageHeader types and Serializers.
+ *
  * @author David Turanski
  * @author Gary Russell
+ * @author Artem Bilan
+ *
  * @since 4.2
  */
 public class MessageKryoRegistrar extends AbstractKryoRegistrar {
 
-	private volatile int messageHeadersRegistrationId = RegistrationIds.DEFAULT_MESSAGEHEADERS_ID;
+	private int genericMessageRegistrationId = RegistrationIds.DEFAULT_GENERIC_MESSAGE_ID;
+
+	private int errorMessageRegistrationId = RegistrationIds.DEFAULT_ERROR_MESSAGE_ID;
+
+	private int adviceMessageRegistrationId = RegistrationIds.DEFAULT_ADVICE_MESSAGE_ID;
+
+	private int mutableMessageRegistrationId = RegistrationIds.DEFAULT_MUTABLE_MESSAGE_ID;
 
-	private volatile int mutableMessageHeadersRegistrationId = RegistrationIds.DEFAULT_MUTABLE_MESSAGEHEADERS_ID;
+	private int messageHeadersRegistrationId = RegistrationIds.DEFAULT_MESSAGEHEADERS_ID;
+
+	private int mutableMessageHeadersRegistrationId = RegistrationIds.DEFAULT_MUTABLE_MESSAGEHEADERS_ID;
+
+	private int hashMapRegistrationId = RegistrationIds.DEFAULT_HASH_MAP_ID;
+
+	private int uuidRegistrationId = RegistrationIds.DEFAULT_UUID_ID;
 
 	/**
-	 * Set the registration id for {@code MessageHeaders}.
+	 * Set the registration id for {@link MessageHeaders}.
 	 * @param messageHeadersRegistrationId the id, default 41.
 	 */
 	public void setMessageHeadersRegistrationId(int messageHeadersRegistrationId) {
 		this.messageHeadersRegistrationId = messageHeadersRegistrationId;
 	}
 
 	/**
-	 * Set the registration id for {@code MutableMessageHeaders}.
+	 * Set the registration id for {@link MutableMessageHeaders}.
 	 * @param mutableMessageHeadersRegistrationId the id, default 42.
 	 */
 	public void setMutableMessageHeadersRegistrationId(int mutableMessageHeadersRegistrationId) {
 		this.mutableMessageHeadersRegistrationId = mutableMessageHeadersRegistrationId;
 	}
 
+	/**
+	 * Set the registration id for {@link GenericMessage}.
+	 * @param genericMessageRegistrationId the id, default 43.
+	 * @since 4.3.23
+	 */
+	public void setGenericMessageRegistrationId(int genericMessageRegistrationId) {
+		this.genericMessageRegistrationId = genericMessageRegistrationId;
+	}
+
+	/**
+	 * Set the registration id for {@link ErrorMessage}.
+	 * @param errorMessageRegistrationId the id, default 44.
+	 * @since 4.3.23
+	 */
+	public void setErrorMessageRegistrationId(int errorMessageRegistrationId) {
+		this.errorMessageRegistrationId = errorMessageRegistrationId;
+	}
+
+	/**
+	 * Set the registration id for {@link AdviceMessage}.
+	 * @param adviceMessageRegistrationId the id, default 45.
+	 * @since 4.3.23
+	 */
+	public void setAdviceMessageRegistrationId(int adviceMessageRegistrationId) {
+		this.adviceMessageRegistrationId = adviceMessageRegistrationId;
+	}
+
+	/**
+	 * Set the registration id for {@link MutableMessage}.
+	 * @param mutableMessageRegistrationId the id, default 46.
+	 * @since 4.3.23
+	 */
+	public void setMutableMessageRegistrationId(int mutableMessageRegistrationId) {
+		this.mutableMessageRegistrationId = mutableMessageRegistrationId;
+	}
+
+	/**
+	 * Set the registration id for {@link HashMap}.
+	 * @param hashMapRegistrationId the id, default 47.
+	 * @since 4.3.23
+	 */
+	public void setHashMapRegistrationId(int hashMapRegistrationId) {
+		this.hashMapRegistrationId = hashMapRegistrationId;
+	}
+
+	/**
+	 * Set the registration id for {@link UUID}.
+	 * @param uuidRegistrationId the id, default 48.
+	 * @since 4.3.23
+	 */
+	public void setUuidRegistrationId(int uuidRegistrationId) {
+		this.uuidRegistrationId = uuidRegistrationId;
+	}
+
+	@Override
+	public void registerTypes(Kryo kryo) {
+		super.registerTypes(kryo);
+		kryo.register(GenericMessage.class, this.genericMessageRegistrationId);
+		kryo.register(ErrorMessage.class, this.errorMessageRegistrationId);
+		kryo.register(AdviceMessage.class, this.adviceMessageRegistrationId);
+		kryo.register(MutableMessage.class, this.mutableMessageRegistrationId);
+		kryo.register(HashMap.class, this.hashMapRegistrationId);
+		kryo.register(UUID.class, this.uuidRegistrationId);
+	}
 
 	@Override
 	public List<Registration> getRegistrations() {

spring-integration-core/src/main/java/org/springframework/integration/codec/kryo/RegistrationIds.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2015-2019 the original author or authors.
+ * Copyright 2015-2020 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -33,6 +33,18 @@ public final class RegistrationIds {
 
 	public static final int DEFAULT_MUTABLE_MESSAGEHEADERS_ID = 42;
 
+	public static final int DEFAULT_GENERIC_MESSAGE_ID = 43;
+
+	public static final int DEFAULT_ERROR_MESSAGE_ID = 44;
+
+	public static final int DEFAULT_ADVICE_MESSAGE_ID = 45;
+
+	public static final int DEFAULT_MUTABLE_MESSAGE_ID = 46;
+
+	public static final int DEFAULT_HASH_MAP_ID = 47;
+
+	public static final int DEFAULT_UUID_ID = 48;
+
 	private RegistrationIds() { }
 
 }