Simplify shared resource setup with Memoized

2 files changed, 6 insertions, 21 deletions

diff --git a/container-core/src/main/java/com/yahoo/container/core/documentapi/DocumentAccessProvider.java b/container-core/src/main/java/com/yahoo/container/core/documentapi/DocumentAccessProvider.java
index be8ba669ec0..54c263e2979 100644
--- a/container-core/src/main/java/com/yahoo/container/core/documentapi/DocumentAccessProvider.java
+++ b/container-core/src/main/java/com/yahoo/container/core/documentapi/DocumentAccessProvider.java
@@ -10,7 +10,7 @@ import com.yahoo.vespa.config.content.DistributionConfig;
 import com.yahoo.vespa.config.content.LoadTypeConfig;
 
 /**
- * Lets a lazily initialised DocumentAccess forwarding to a real MessageBusDocumentAccess be injected in containers.
+ * Lets a lazily initialised DocumentAccess that forwards to a MessageBusDocumentAccess be injected in containers.
  *
  * @author jonmv
  */
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 7c5bebc47e8..7fc11c6bbd8 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,6 +20,7 @@ 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;
@@ -37,8 +38,7 @@ public class VespaDocumentAccess extends DocumentAccess {
 
     private final MessageBusParams parameters;
 
-    private final AtomicReference<DocumentAccess> delegate = new AtomicReference<>();
-    private boolean shutDown = false;
+    private final Memoized<DocumentAccess, RuntimeException> delegate;
 
     VespaDocumentAccess(DocumentmanagerConfig documentmanagerConfig,
                         String slobroksConfigId,
@@ -51,19 +51,11 @@ 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() {
-        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);
-        });
+        return delegate.get();
     }
 
     @Override
@@ -72,14 +64,7 @@ public class VespaDocumentAccess extends DocumentAccess {
     }
 
     void protectedShutdown() {
-        delegate.updateAndGet(access -> {
-            super.shutdown();
-            shutDown = true;
-            if (access != null)
-                access.shutdown();
-
-            return null;
-        });
+        delegate.close();
     }
 
     @Override