Use an injected VespaDocumentAccess for streaming search

1 files changed, 19 insertions, 15 deletions

diff --git a/container-core/src/main/java/com/yahoo/container/core/documentapi/VespaDocumentAccess.java b/container-core/src/main/java/com/yahoo/container/core/documentapi/VespaDocumentAccess.java
index d55df15b2fd..bb78ee3022f 100644
--- a/container-core/src/main/java/com/yahoo/container/core/documentapi/VespaDocumentAccess.java
+++ b/container-core/src/main/java/com/yahoo/container/core/documentapi/VespaDocumentAccess.java
@@ -24,6 +24,8 @@ import com.yahoo.documentapi.messagebus.protocol.DocumentProtocolPoliciesConfig;
 import com.yahoo.vespa.config.content.DistributionConfig;
 import com.yahoo.vespa.config.content.LoadTypeConfig;
 
+import java.util.concurrent.atomic.AtomicReference;
+
 /**
  * Wraps a lazily initialised MessageBusDocumentAccess. Lazy to allow it to always be set up.
  * Inject this class directly (instead of DocumentAccess) for use in internal code.
@@ -33,9 +35,8 @@ import com.yahoo.vespa.config.content.LoadTypeConfig;
 public class VespaDocumentAccess extends DocumentAccess {
 
     private final MessageBusParams parameters;
-    private final Object monitor = new Object();
 
-    private DocumentAccess delegate = null;
+    private final AtomicReference<MessageBusDocumentAccess> delegate = new AtomicReference<>();
     private boolean shutDown = false;
 
     VespaDocumentAccess(DocumentmanagerConfig documentmanagerConfig,
@@ -52,26 +53,29 @@ public class VespaDocumentAccess extends DocumentAccess {
         this.parameters.getMessageBusParams().setMessageBusConfig(messagebusConfig);
     }
 
-    private DocumentAccess delegate() {
-        synchronized (monitor) {
-            if (delegate == null) {
-                if (shutDown)
-                    throw new IllegalStateException("This document access has been shut down");
+    public MessageBusDocumentAccess delegate() {
+        MessageBusDocumentAccess access = delegate.getAcquire();
+        return access != null ? access : delegate.updateAndGet(value -> {
+            if (value != null)
+                return value;
+
+            if (shutDown)
+                throw new IllegalStateException("This document access has been shut down");
 
-                delegate = new MessageBusDocumentAccess(parameters);
-            }
-            return delegate;
-        }
+            return new MessageBusDocumentAccess(parameters);
+        });
     }
 
     @Override
     public void shutdown() {
-        synchronized (monitor) {
+        delegate.updateAndGet(access -> {
             super.shutdown();
             shutDown = true;
-            if (delegate != null)
-                delegate.shutdown();
-        }
+            if (access != null)
+                access.shutdown();
+
+            return null;
+        });
     }
 
     @Override