Fix ThreadLocal Issue with Repository Save. (#1840)

The issue was introduced when the Mono.deferContextual() was added to
determine if the save() is in a transaction. It may be executing in a
different thread when the PseudoArgs (scope, collection, and options)
are retrieved ThreadLocal. This change ensures scope and collection
are retrieved, but options are ignored and discarded.

Closes #1838.

src/main/java/org/springframework/data/couchbase/core/ReactiveCouchbaseTemplate.java

@@ -53,7 +53,7 @@ public class ReactiveCouchbaseTemplate implements ReactiveCouchbaseOperations, A
  private final CouchbaseConverter converter;
  private final PersistenceExceptionTranslator exceptionTranslator;
  private final ReactiveCouchbaseTemplateSupport templateSupport;
- private ThreadLocal<PseudoArgs<?>> threadLocalArgs = new ThreadLocal<>();
+ private final ThreadLocal<PseudoArgs<?>> threadLocalArgs = new ThreadLocal<>();
  private final QueryScanConsistency scanConsistency;
 
  public ReactiveCouchbaseTemplate(final CouchbaseClientFactory clientFactory, final CouchbaseConverter converter) {
@@ -257,14 +257,6 @@ public PseudoArgs<?> getPseudoArgs() {
  * set the ThreadLocal field
  */
  public void setPseudoArgs(PseudoArgs<?> threadLocalArgs) {
- if (this.threadLocalArgs == null) {
- synchronized (this) {
- if (this.threadLocalArgs == null) {
- this.threadLocalArgs = new ThreadLocal<>();
- }
- }
- }
-
  this.threadLocalArgs.set(threadLocalArgs);
  }
 

src/main/java/org/springframework/data/couchbase/repository/support/CouchbaseRepositoryBase.java

@@ -18,7 +18,9 @@
 
 import java.lang.reflect.AnnotatedElement;
 
+import org.springframework.data.couchbase.core.ReactiveCouchbaseTemplate;
 import org.springframework.data.couchbase.core.query.OptionsBuilder;
+import org.springframework.data.couchbase.core.support.PseudoArgs;
 import org.springframework.data.couchbase.repository.Collection;
 import org.springframework.data.couchbase.repository.ScanConsistency;
 import org.springframework.data.couchbase.repository.Scope;
@@ -35,7 +37,7 @@
  *
  * @author Michael Reiche
  */
-public class CouchbaseRepositoryBase<T, ID> {
+public abstract class CouchbaseRepositoryBase<T, ID> {
 
  /**
  * Contains information about the entity being used in this repository.
@@ -82,9 +84,11 @@ <S extends T> String getId(S entity) {
 
  protected String getScope() {
  String fromAnnotation = OptionsBuilder.annotationString(Scope.class, CollectionIdentifier.DEFAULT_SCOPE,
- new AnnotatedElement[] { getJavaType(), repositoryInterface });
+ new AnnotatedElement[] { getJavaType(), getRepositoryInterface() });
  String fromMetadata = crudMethodMetadata != null ? crudMethodMetadata.getScope() : null;
- return OptionsBuilder.fromFirst(CollectionIdentifier.DEFAULT_SCOPE, fromMetadata, fromAnnotation);
+ PseudoArgs<?> pa = getReactiveTemplate().getPseudoArgs();
+ String fromThreadLocal = pa != null ? pa.getScope() : null;
+ return OptionsBuilder.fromFirst(CollectionIdentifier.DEFAULT_SCOPE, fromThreadLocal, fromMetadata, fromAnnotation);
  }
 
  /**
@@ -96,12 +100,18 @@ protected String getScope() {
  * 1. repository.withCollection()
  */
  protected String getCollection() {
- String fromAnnotation = OptionsBuilder.annotationString(Collection.class, CollectionIdentifier.DEFAULT_COLLECTION,
- new AnnotatedElement[] { getJavaType(), repositoryInterface });
+ String fromAnnotation = OptionsBuilder.annotationString(Collection.class,
+ CollectionIdentifier.DEFAULT_COLLECTION,
+ new AnnotatedElement[] { getJavaType(), getRepositoryInterface() });
  String fromMetadata = crudMethodMetadata != null ? crudMethodMetadata.getCollection() : null;
- return OptionsBuilder.fromFirst(CollectionIdentifier.DEFAULT_COLLECTION, fromMetadata, fromAnnotation);
+ PseudoArgs<?> pa = getReactiveTemplate().getPseudoArgs();
+ String fromThreadLocal = pa != null ? pa.getCollection() : null;
+ return OptionsBuilder.fromFirst(CollectionIdentifier.DEFAULT_COLLECTION, fromThreadLocal, fromMetadata,
+ fromAnnotation);
  }
 
+ protected abstract ReactiveCouchbaseTemplate getReactiveTemplate();
+
  /**
  * Get the QueryScanConsistency from <br>
  * 1. The method annotation (method *could* be available from crudMethodMetadata)<br>
@@ -132,4 +142,5 @@ QueryScanConsistency getQueryScanConsistency() {
  void setRepositoryMethodMetadata(CrudMethodMetadata crudMethodMetadata) {
  this.crudMethodMetadata = crudMethodMetadata;
  }
+
 }

src/main/java/org/springframework/data/couchbase/repository/support/SimpleCouchbaseRepository.java

@@ -23,8 +23,8 @@
 import java.util.stream.Collectors;
 
 import org.springframework.data.couchbase.core.CouchbaseOperations;
-import org.springframework.data.couchbase.core.mapping.CouchbasePersistentEntity;
-import org.springframework.data.couchbase.core.mapping.CouchbasePersistentProperty;
+import org.springframework.data.couchbase.core.CouchbaseTemplate;
+import org.springframework.data.couchbase.core.ReactiveCouchbaseTemplate;
 import org.springframework.data.couchbase.core.query.Query;
 import org.springframework.data.couchbase.repository.CouchbaseRepository;
 import org.springframework.data.couchbase.repository.query.CouchbaseEntityInformation;
@@ -37,7 +37,6 @@
 import org.springframework.util.Assert;
 
 import com.couchbase.client.java.query.QueryScanConsistency;
-import org.springframework.util.ReflectionUtils;
 
 /**
  * Repository base implementation for Couchbase.
@@ -71,7 +70,13 @@ public SimpleCouchbaseRepository(CouchbaseEntityInformation<T, String> entityInf
  @Override
  @SuppressWarnings("unchecked")
  public <S extends T> S save(S entity) {
- return operations.save(entity, getScope(), getCollection());
+ String scopeName = getScope();
+ String collectionName = getCollection();
+ // clear out the PseudoArgs here as whatever is called by operations.save() could be in a different thread.
+ // not that this will also clear out Options, but that's ok as any options would not work
+ // with all of insert/upsert/replace. If Options are needed, use template.insertById/upsertById/replaceById
+ getReactiveTemplate().setPseudoArgs(null);
+ return operations.save(entity, scopeName, collectionName);
  }
 
  @Override
@@ -177,4 +182,8 @@ public CouchbaseOperations getOperations() {
  return operations;
  }
 
+ @Override
+ protected ReactiveCouchbaseTemplate getReactiveTemplate() {
+ return ((CouchbaseTemplate) getOperations()).reactive();
+ }
 }

src/main/java/org/springframework/data/couchbase/repository/support/SimpleReactiveCouchbaseRepository.java

@@ -26,15 +26,13 @@
 import org.reactivestreams.Publisher;
 import org.springframework.data.couchbase.core.CouchbaseOperations;
 import org.springframework.data.couchbase.core.ReactiveCouchbaseOperations;
-import org.springframework.data.couchbase.core.mapping.CouchbasePersistentEntity;
-import org.springframework.data.couchbase.core.mapping.CouchbasePersistentProperty;
+import org.springframework.data.couchbase.core.ReactiveCouchbaseTemplate;
 import org.springframework.data.couchbase.core.query.Query;
 import org.springframework.data.couchbase.repository.ReactiveCouchbaseRepository;
 import org.springframework.data.couchbase.repository.query.CouchbaseEntityInformation;
 import org.springframework.data.domain.Sort;
 import org.springframework.data.util.Streamable;
 import org.springframework.util.Assert;
-import org.springframework.util.ReflectionUtils;
 
 /**
  * Reactive repository base implementation for Couchbase.
@@ -76,7 +74,13 @@ public Flux<T> findAll(Sort sort) {
  @SuppressWarnings("unchecked")
  @Override
  public <S extends T> Mono<S> save(S entity) {
- return save(entity, getScope(), getCollection());
+ String scopeName = getScope();
+ String collectionName = getCollection();
+ // clear out the PseudoArgs here as whatever is called by operations.save() could be in a different thread.
+ // not that this will also clear out Options, but that's ok as any options would not work
+ // with all of insert/upsert/replace. If Options are needed, use template.insertById/upsertById/replaceById
+ getReactiveTemplate().setPseudoArgs(null);
+ return operations.save(entity, scopeName, collectionName);
  }
 
  @Override
@@ -227,4 +231,9 @@ public ReactiveCouchbaseOperations getOperations() {
  return operations;
  }
 
+ @Override
+ protected ReactiveCouchbaseTemplate getReactiveTemplate() {
+ return (ReactiveCouchbaseTemplate) getOperations();
+ }
+
 }

src/test/java/org/springframework/data/couchbase/domain/ReactiveAirportMustScopeRepository.java

@@ -0,0 +1,9 @@
+package org.springframework.data.couchbase.domain;
+
+import org.springframework.data.couchbase.repository.Collection;
+import org.springframework.data.couchbase.repository.Scope;
+
+@Scope("must set scope name")
+@Collection("my_collection")
+public interface ReactiveAirportMustScopeRepository extends ReactiveAirportRepository {
+}

src/test/java/org/springframework/data/couchbase/repository/query/ReactiveCouchbaseRepositoryQueryCollectionIntegrationTests.java

@@ -18,7 +18,11 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
+import reactor.core.Disposable;
+
 import java.util.List;
+import java.util.Random;
+import java.util.UUID;
 
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.AfterEach;
@@ -31,6 +35,7 @@
 import org.springframework.data.couchbase.core.ReactiveCouchbaseTemplate;
 import org.springframework.data.couchbase.domain.Airport;
 import org.springframework.data.couchbase.domain.ConfigScoped;
+import org.springframework.data.couchbase.domain.ReactiveAirportMustScopeRepository;
 import org.springframework.data.couchbase.domain.ReactiveAirportRepository;
 import org.springframework.data.couchbase.domain.ReactiveAirportRepositoryAnnotated;
 import org.springframework.data.couchbase.domain.ReactiveUserColRepository;
@@ -61,6 +66,7 @@ public class ReactiveCouchbaseRepositoryQueryCollectionIntegrationTests extends
 
  @Autowired ReactiveAirportRepository reactiveAirportRepository;
  @Autowired ReactiveAirportRepositoryAnnotated reactiveAirportRepositoryAnnotated;
+ @Autowired ReactiveAirportMustScopeRepository reactiveAirportMustScopeRepository;
  @Autowired ReactiveUserColRepository userColRepository;
  @Autowired public CouchbaseTemplate couchbaseTemplate;
  @Autowired public ReactiveCouchbaseTemplate reactiveCouchbaseTemplate;
@@ -116,6 +122,21 @@ public void myTest() {
 
  }
 
+ @Test
+ void testThreadLocal() throws InterruptedException {
+
+ String scopeName = "my_scope";
+ String id = UUID.randomUUID().toString();
+
+ Airport airport = new Airport(id, "testThreadLocal", "icao");
+ Disposable s = reactiveAirportMustScopeRepository.withScope(scopeName).findById(airport.getId()).doOnNext(u -> {
+ throw new RuntimeException("User already Exists! " + u);
+ }).then(reactiveAirportMustScopeRepository.withScope(scopeName).save(airport))
+ .subscribe(u -> LOGGER.info("User Persisted Successfully! {}", u));
+
+ reactiveAirportMustScopeRepository.withScope(scopeName).deleteById(id).block();
+ }
+
  /**
  * can test against _default._default without setting up additional scope/collection and also test for collections and
  * scopes that do not exist These same tests should be repeated on non-default scope and collection in a test that