MessageHeaders serialization preserves non-serializable entries in original instance

Issue: SPR-13541

spring-messaging/src/main/java/org/springframework/messaging/MessageHeaders.java

@@ -20,11 +20,10 @@ import java.io.IOException;
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
 import java.io.Serializable;
-import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.List;
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.UUID;
@@ -64,6 +63,7 @@ import org.springframework.util.IdGenerator;
  * @author Arjen Poutsma
  * @author Mark Fisher
  * @author Gary Russell
+ * @author Juergen Hoeller
  * @since 4.0
  * @see org.springframework.messaging.support.MessageBuilder
  * @see org.springframework.messaging.support.MessageHeaderAccessor
@@ -140,6 +140,21 @@ public class MessageHeaders implements Map<String, Object>, Serializable {
 		}
 	}
 
+	/**
+	 * Copy constructor which allows for ignoring certain entries.
+	 * Used for serialization without non-serializable entries.
+	 * @param original the MessageHeaders to copy
+	 * @param keysToIgnore the keys of the entries to ignore
+	 */
+	private MessageHeaders(MessageHeaders original, Set<String> keysToIgnore) {
+		this.headers = new HashMap<String, Object>(original.headers.size() - keysToIgnore.size());
+		for (Map.Entry<String, Object> entry : original.headers.entrySet()) {
+			if (!keysToIgnore.contains(entry.getKey())) {
+				this.headers.put(entry.getKey(), entry.getValue());
+			}
+		}
+	}
+
 
 	protected Map<String, Object> getRawHeaders() {
 		return this.headers;
@@ -165,6 +180,7 @@ public class MessageHeaders implements Map<String, Object>, Serializable {
 		return get(ERROR_CHANNEL);
 	}
 
+
 	@SuppressWarnings("unchecked")
 	public <T> T get(Object key, Class<T> type) {
 		Object value = this.headers.get(key);
@@ -179,23 +195,6 @@ public class MessageHeaders implements Map<String, Object>, Serializable {
 	}
 
 
-	@Override
-	public boolean equals(Object other) {
-		return (this == other ||
-				(other instanceof MessageHeaders && this.headers.equals(((MessageHeaders) other).headers)));
-	}
-
-	@Override
-	public int hashCode() {
-		return this.headers.hashCode();
-	}
-
-	@Override
-	public String toString() {
-		return this.headers.toString();
-	}
-
-
 	// Delegating Map implementation
 
 	public boolean containsKey(Object key) {
@@ -269,23 +268,47 @@ public class MessageHeaders implements Map<String, Object>, Serializable {
 	// Serialization methods
 
 	private void writeObject(ObjectOutputStream out) throws IOException {
-		List<String> keysToRemove = new ArrayList<String>();
+		Set<String> keysToIgnore = new HashSet<String>();
 		for (Map.Entry<String, Object> entry : this.headers.entrySet()) {
 			if (!(entry.getValue() instanceof Serializable)) {
-				keysToRemove.add(entry.getKey());
+				keysToIgnore.add(entry.getKey());
 			}
 		}
-		for (String key : keysToRemove) {
-			if (logger.isInfoEnabled()) {
-				logger.info("Removing non-serializable header: " + key);
-			}
-			this.headers.remove(key);
+
+		if (keysToIgnore.isEmpty()) {
+			// All entries are serializable -> serialize the regular MessageHeaders instance
+			out.defaultWriteObject();
+		}
+		else {
+			// Some non-serializable entries -> serialize a temporary MessageHeaders copy
+			if (logger.isDebugEnabled()) {
+				logger.debug("Ignoring non-serializable message headers: " + keysToIgnore);
+			}
+			out.writeObject(new MessageHeaders(this, keysToIgnore));
 		}
-		out.defaultWriteObject();
 	}
 
 	private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
 		in.defaultReadObject();
 	}
 
+
+	// equals, hashCode, toString
+
+	@Override
+	public boolean equals(Object other) {
+		return (this == other ||
+				(other instanceof MessageHeaders && this.headers.equals(((MessageHeaders) other).headers)));
+	}
+
+	@Override
+	public int hashCode() {
+		return this.headers.hashCode();
+	}
+
+	@Override
+	public String toString() {
+		return this.headers.toString();
+	}
+
 }