DATAMONGO-1005 - Improve cycle-detection for DbRef's.

Introduced ObjectPath that collects the target objects while converting a DBObject to a domain object. If we detect that a potentially nested DBRef points to an object that is already under construction we simply return a reference to that object in order to avoid StackOverFlowErrors.

Original pull request: #209.

Author: Thomas Darimont

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/MappingMongoConverter.java

@@ -20,6 +20,7 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -194,6 +195,8 @@ protected <S extends Object> S read(TypeInformation<S> type, DBObject dbo, Objec
 			return null;
 		}
 
+		Object parentObject = ObjectPath.unwrapCurrentFromPotentialPath(parent);
+
 		TypeInformation<? extends S> typeToUse = typeMapper.readType(dbo, type);
 		Class<? extends S> rawType = typeToUse.getType();
 
@@ -206,11 +209,11 @@ protected <S extends Object> S read(TypeInformation<S> type, DBObject dbo, Objec
 		}
 
 		if (typeToUse.isCollectionLike() && dbo instanceof BasicDBList) {
-			return (S) readCollectionOrArray(typeToUse, (BasicDBList) dbo, parent);
+			return (S) readCollectionOrArray(typeToUse, (BasicDBList) dbo, parentObject);
 		}
 
 		if (typeToUse.isMap()) {
-			return (S) readMap(typeToUse, dbo, parent);
+			return (S) readMap(typeToUse, dbo, parentObject);
 		}
 
 		// Retrieve persistent entity info
@@ -238,22 +241,37 @@ private <S extends Object> S read(final MongoPersistentEntity<S> entity, final D
 
 		final DefaultSpELExpressionEvaluator evaluator = new DefaultSpELExpressionEvaluator(dbo, spELContext);
 
-		ParameterValueProvider<MongoPersistentProperty> provider = getParameterProvider(entity, dbo, evaluator, parent);
+		final ObjectPath parentPath = ObjectPath.toPath(parent);
+
+		ParameterValueProvider<MongoPersistentProperty> provider = getParameterProvider(entity, dbo, evaluator,
+				parentPath.getCurrentObject());
 		EntityInstantiator instantiator = instantiators.getInstantiatorFor(entity);
 		S instance = instantiator.createInstance(entity, provider);
 
 		final BeanWrapper<S> wrapper = BeanWrapper.create(instance, conversionService);
 		final S result = wrapper.getBean();
 
+		// make sure id property is set before all other properties
+		if (entity.getIdProperty() != null) {
+			wrapper.setProperty(entity.getIdProperty(), getValueInternal(entity.getIdProperty(), dbo, evaluator, result));
+		}
+
+		final ObjectPath currentPath = parentPath.push(result);
+
 		// Set properties not already set in the constructor
 		entity.doWithProperties(new PropertyHandler<MongoPersistentProperty>() {
 			public void doWithPersistentProperty(MongoPersistentProperty prop) {
 
+				// we skip the id property since it was already set
+				if (entity.getIdProperty() != null && entity.getIdProperty().equals(prop)) {
+					return;
+				}
+
 				if (!dbo.containsField(prop.getFieldName()) || entity.isConstructorArgument(prop)) {
 					return;
 				}
 
-				wrapper.setProperty(prop, getValueInternal(prop, dbo, evaluator, result));
+				wrapper.setProperty(prop, getValueInternal(prop, dbo, evaluator, currentPath.getCurrentObject()));
 			}
 		});
 
@@ -273,7 +291,7 @@ public void doWithAssociation(Association<MongoPersistentProperty> association)
 
 					@Override
 					public Object resolve(MongoPersistentProperty property) {
-						return getValueInternal(property, dbo, evaluator, parent);
+						return getValueInternal(property, dbo, evaluator, currentPath);
 					}
 				}));
 			}
@@ -1097,21 +1115,75 @@ protected <T> T potentiallyConvertSpelValue(Object object, Parameter<T, MongoPer
 	@SuppressWarnings("unchecked")
 	private <T> T readValue(Object value, TypeInformation<?> type, Object parent) {
 
+		ObjectPath objectStack = ObjectPath.toPath(parent);
+
 		Class<?> rawType = type.getType();
 
 		if (conversions.hasCustomReadTarget(value.getClass(), rawType)) {
 			return (T) conversionService.convert(value, rawType);
 		} else if (value instanceof DBRef) {
-			return (T) (rawType.equals(DBRef.class) ? value : read(type, readRef((DBRef) value), parent));
+			return potentiallyReadOrResolveDbRef((DBRef) value, type, objectStack, rawType);
 		} else if (value instanceof BasicDBList) {
-			return (T) readCollectionOrArray(type, (BasicDBList) value, parent);
+			return (T) readCollectionOrArray(type, (BasicDBList) value, objectStack.getCurrentObject());
 		} else if (value instanceof DBObject) {
-			return (T) read(type, (DBObject) value, parent);
+			return (T) read(type, (DBObject) value, objectStack.getCurrentObject());
 		} else {
 			return (T) getPotentiallyConvertedSimpleRead(value, rawType);
 		}
 	}
 
+	@SuppressWarnings("unchecked")
+	private <T> T potentiallyReadOrResolveDbRef(DBRef dbref, TypeInformation<?> type, ObjectPath path, Class<?> rawType) {
+
+		if (rawType.equals(DBRef.class)) {
+			return (T) dbref;
+		}
+
+		Object object = getObjectFromPathForRefOrNull(path, dbref);
+
+		if (object != null) {
+			return (T) object;
+		}
+
+		return (T) (object != null ? object : read(type, readRef(dbref), path.getCurrentObject()));
+	}
+
+	/**
+	 * Returns the object from the given {@link ObjectPath} iff the given {@link DBRef} points to it or {@literal null}.
+	 * 
+	 * @param path
+	 * @param dbref
+	 * @return
+	 */
+	private Object getObjectFromPathForRefOrNull(ObjectPath path, DBRef dbref) {
+
+		if (path == null || dbref == null) {
+			return null;
+		}
+
+		for (Object object : path) {
+
+			if (object == null) {
+				continue;
+			}
+
+			MongoPersistentEntity<?> pe = getMappingContext().getPersistentEntity(object.getClass());
+
+			if (pe == null || pe.getIdProperty() == null) {
+				continue;
+			}
+
+			BeanWrapper<Object> parentBeanWrapper = BeanWrapper.create(object, conversionService);
+			Object parentId = parentBeanWrapper.getProperty(pe.getIdProperty());
+
+			if (dbref.getRef().equals(pe.getCollection()) && dbref.getId().equals(parentId)) {
+				return object;
+			}
+		}
+
+		return null;
+	}
+
 	/**
 	 * Performs the fetch operation for the given {@link DBRef}.
 	 * 
@@ -1121,4 +1193,70 @@ private <T> T readValue(Object value, TypeInformation<?> type, Object parent) {
 	DBObject readRef(DBRef ref) {
 		return ref.fetch();
 	}
+
+	/**
+	 * An immutable ordered set of target objects for {@link DBObject} to {@link Object} conversions. Object paths can be
+	 * constructed by the {@link #toPath(Object)} method and extended via {@link #push(Object)}.
+	 * 
+	 * @author Thomas Darimont
+	 * @since 1.6
+	 */
+	static class ObjectPath implements Iterable<Object> {
+
+		private final List<Object> stack;
+
+		private ObjectPath(Object current) {
+			this(null, current);
+		}
+
+		private ObjectPath(ObjectPath parent, Object current) {
+
+			if (parent == null) {
+				this.stack = Collections.singletonList(current);
+			} else {
+				List<Object> list = new ArrayList<Object>(parent.stack);
+				list.add(current);
+				this.stack = list;
+			}
+		}
+
+		/**
+		 * Returns a copy of the {@link ObjectPath} with the given {@link Object} as current object.
+		 * 
+		 * @param o
+		 * @return
+		 */
+		public ObjectPath push(Object o) {
+			return new ObjectPath(this, o);
+		}
+
+		public Object getRootObject() {
+			return stack.get(0);
+		}
+
+		public Object getCurrentObject() {
+			return stack.get(stack.size() - 1);
+		}
+
+		@Override
+		public Iterator<Object> iterator() {
+			return stack.iterator();
+		}
+
+		/**
+		 * Returns the {@link ObjectPath} represented by the given {@link Object} or creates a new {@link ObjectPath}
+		 * wrapping the given {@link Object}.
+		 * 
+		 * @param object
+		 * @return
+		 */
+		public static ObjectPath toPath(Object object) {
+			return object instanceof ObjectPath ? ((ObjectPath) object) : new ObjectPath(object);
+		}
+
+		public static Object unwrapCurrentFromPotentialPath(Object potentialObjectPath) {
+			return potentialObjectPath instanceof ObjectPath ? ((ObjectPath) potentialObjectPath).getCurrentObject()
+					: potentialObjectPath;
+		}
+	}
 }

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/convert/DbRefMappingMongoConverterUnitTests.java

@@ -15,7 +15,7 @@
  */
 package org.springframework.data.mongodb.core.convert;
 
-import static org.hamcrest.Matchers.*;
+import static org.hamcrest.CoreMatchers.*;
 import static org.junit.Assert.*;
 import static org.mockito.Matchers.*;
 import static org.mockito.Mockito.*;
@@ -450,6 +450,47 @@ public void shouldNotGenerateLazyLoadingProxyForNullValues() {
 		assertThat(result.dbRefToConcreteTypeWithPersistenceConstructorWithoutDefaultConstructor, is(nullValue()));
 	}
 
+	/**
+	 * @see DATAMONGO-1005
+	 */
+	@Test
+	public void shouldBeAbleToStoreDirectReferencesToSelf() {
+
+		DBObject dbo = new BasicDBObject();
+
+		ClassWithDbRefField o = new ClassWithDbRefField();
+		o.id = "123";
+		o.reference = o;
+		converter.write(o, dbo);
+
+		ClassWithDbRefField found = converter.read(ClassWithDbRefField.class, dbo);
+
+		assertThat(found, is(notNullValue()));
+		assertThat(found.reference, is(found));
+	}
+
+	/**
+	 * @see DATAMONGO-1005
+	 */
+	@Test
+	public void shouldBeAbleToStoreNestedReferencesToSelf() {
+
+		DBObject dbo = new BasicDBObject();
+
+		ClassWithNestedDbRefField o = new ClassWithNestedDbRefField();
+		o.id = "123";
+		o.nested = new NestedReferenceHolder();
+		o.nested.reference = o;
+
+		converter.write(o, dbo);
+
+		ClassWithNestedDbRefField found = converter.read(ClassWithNestedDbRefField.class, dbo);
+
+		assertThat(found, is(notNullValue()));
+		assertThat(found.nested, is(notNullValue()));
+		assertThat(found.nested.reference, is(found));
+	}
+
 	private Object transport(Object result) {
 		return SerializationUtils.deserialize(SerializationUtils.serialize(result));
 	}
@@ -626,4 +667,22 @@ static class WithObjectMethodOverrideLazyDbRefs {
 		@org.springframework.data.mongodb.core.mapping.DBRef(lazy = true) EqualsAndHashCodeObjectMethodOverrideLazyDbRefTarget dbRefEqualsAndHashcodeObjectMethodOverride2;
 		@org.springframework.data.mongodb.core.mapping.DBRef(lazy = true) EqualsAndHashCodeObjectMethodOverrideLazyDbRefTarget dbRefEqualsAndHashcodeObjectMethodOverride1;
 	}
+
+	class ClassWithDbRefField {
+
+		String id;
+		@org.springframework.data.mongodb.core.mapping.DBRef ClassWithDbRefField reference;
+	}
+
+	static class NestedReferenceHolder {
+
+		String id;
+		@org.springframework.data.mongodb.core.mapping.DBRef ClassWithNestedDbRefField reference;
+	}
+
+	static class ClassWithNestedDbRefField {
+
+		String id;
+		NestedReferenceHolder nested;
+	}
 }