1 files changed, 21 insertions, 5 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 c708e87e6c5..7c5bebc47e8 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
@@ -20,8 +20,8 @@ import com.yahoo.documentapi.messagebus.MessageBusParams;
 import com.yahoo.documentapi.messagebus.protocol.DocumentProtocolPoliciesConfig;
 import com.yahoo.messagebus.MessagebusConfig;
 import com.yahoo.vespa.config.content.DistributionConfig;
-import com.yahoo.yolean.concurrent.Memoized;
 
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -37,7 +37,8 @@ public class VespaDocumentAccess extends DocumentAccess {
 
     private final MessageBusParams parameters;
 
-    private final Memoized<DocumentAccess, RuntimeException> delegate;
+    private final AtomicReference<DocumentAccess> delegate = new AtomicReference<>();
+    private boolean shutDown = false;
 
     VespaDocumentAccess(DocumentmanagerConfig documentmanagerConfig,
                         String slobroksConfigId,
@@ -50,11 +51,19 @@ public class VespaDocumentAccess extends DocumentAccess {
         this.parameters.setDocumentmanagerConfig(documentmanagerConfig);
         this.parameters.getRPCNetworkParams().setSlobrokConfigId(slobroksConfigId);
         this.parameters.getMessageBusParams().setMessageBusConfig(messagebusConfig);
-        this.delegate = new Memoized<>(() -> new MessageBusDocumentAccess(parameters), DocumentAccess::shutdown);
     }
 
     public DocumentAccess delegate() {
-        return delegate.get();
+        DocumentAccess 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");
+
+            return new MessageBusDocumentAccess(parameters);
+        });
     }
 
     @Override
@@ -63,7 +72,14 @@ public class VespaDocumentAccess extends DocumentAccess {
     }
 
     void protectedShutdown() {
-        delegate.close();
+        delegate.updateAndGet(access -> {
+            super.shutdown();
+            shutDown = true;
+            if (access != null)
+                access.shutdown();
+
+            return null;
+        });
     }
 
     @Override