Construct primitive session client asynchronously to allow external s…

…ession ID generation.
atomix · Jun 28, 2018 · b26b60f · b26b60f
1 parent cb2cbac
commit b26b60f
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 33 deletions.
diff --git a/primitive/src/main/java/io/atomix/primitive/session/impl/RecoveringSessionClient.java b/primitive/src/main/java/io/atomix/primitive/session/impl/RecoveringSessionClient.java
@@ -50,7 +50,7 @@ public class RecoveringSessionClient implements SessionClient {
   private final PartitionId partitionId;
   private final String name;
   private final PrimitiveType primitiveType;
-  private final Supplier<SessionClient> proxyFactory;
+  private final Supplier<CompletableFuture<SessionClient>> proxyFactory;
   private final ThreadContext context;
   private Logger log;
   private volatile CompletableFuture<SessionClient> connectFuture;
@@ -67,7 +67,7 @@ public RecoveringSessionClient(
       PartitionId partitionId,
       String name,
       PrimitiveType primitiveType,
-      Supplier<SessionClient> sessionFactory,
+      Supplier<CompletableFuture<SessionClient>> sessionFactory,
       ThreadContext context) {
     this.partitionId = checkNotNull(partitionId);
     this.name = checkNotNull(name);
@@ -160,7 +160,7 @@ private void recover() {
    */
   private void openProxy(CompletableFuture<SessionClient> future) {
     log.debug("Opening proxy session");
-    proxyFactory.get().connect().whenComplete((proxy, error) -> {
+    proxyFactory.get().thenCompose(proxy -> proxy.connect()).whenComplete((proxy, error) -> {
       if (error == null) {
         synchronized (this) {
           this.log = ContextualLoggerFactory.getLogger(getClass(), LoggerContext.builder(SessionClient.class)

diff --git a/protocols/primary-backup/src/main/java/io/atomix/protocols/backup/PrimaryBackupClient.java b/protocols/primary-backup/src/main/java/io/atomix/protocols/backup/PrimaryBackupClient.java
@@ -101,21 +101,22 @@ public PrimaryBackupSessionClient.Builder sessionBuilder(String primitiveName, P
     return new PrimaryBackupSessionClient.Builder() {
       @Override
       public SessionClient build() {
-        Supplier<SessionClient> proxyBuilder = () -> new PrimaryBackupSessionClient(
+        Supplier<CompletableFuture<SessionClient>> proxyBuilder = () -> sessionIdService.nextSessionId()
-            clientName,
+            .thenApply(sessionId -> new PrimaryBackupSessionClient(
-            partitionId,
+                clientName,
-            Futures.get(sessionIdService.nextSessionId()),
+                partitionId,
-            primitiveType,
+                sessionId,
-            new PrimitiveDescriptor(
+                primitiveType,
-                primitiveName,
+                new PrimitiveDescriptor(
-                primitiveType.name(),
+                    primitiveName,
-                configBytes,
+                    primitiveType.name(),
-                numBackups,
+                    configBytes,
-                replication),
+                    numBackups,
-            clusterMembershipService,
+                    replication),
-            PrimaryBackupClient.this.protocol,
+                clusterMembershipService,
-            primaryElection,
+                PrimaryBackupClient.this.protocol,
-            threadContextFactory.createContext());
+                primaryElection,
+                threadContextFactory.createContext()));
 
         SessionClient proxy;
         ThreadContext context = threadContextFactory.createContext();
@@ -128,7 +129,7 @@ public SessionClient build() {
               proxyBuilder,
               context);
         } else {
-          proxy = proxyBuilder.get();
+          proxy = Futures.get(proxyBuilder.get());
         }
 
         // If max retries is set, wrap the client in a retrying proxy client.

diff --git a/protocols/raft/src/main/java/io/atomix/protocols/raft/impl/DefaultRaftClient.java b/protocols/raft/src/main/java/io/atomix/protocols/raft/impl/DefaultRaftClient.java
@@ -129,19 +129,20 @@ public RaftSessionClient.Builder sessionBuilder(String primitiveName, PrimitiveT
       @Override
       public SessionClient build() {
         // Create a proxy builder that uses the session manager to open a session.
-        Supplier<SessionClient> proxyFactory = () -> new DefaultRaftSessionClient(
+        Supplier<CompletableFuture<SessionClient>> proxyFactory = () -> CompletableFuture.completedFuture(
-            primitiveName,
+            new DefaultRaftSessionClient(
-            primitiveType,
+                primitiveName,
-            serviceConfig,
+                primitiveType,
-            partitionId,
+                serviceConfig,
-            DefaultRaftClient.this.protocol,
+                partitionId,
-            selectorManager,
+                DefaultRaftClient.this.protocol,
-            sessionManager,
+                selectorManager,
-            readConsistency,
+                sessionManager,
-            communicationStrategy,
+                readConsistency,
-            threadContextFactory.createContext(),
+                communicationStrategy,
-            minTimeout,
+                threadContextFactory.createContext(),
-            maxTimeout);
+                minTimeout,
+                maxTimeout));
 
         SessionClient proxy;
 
@@ -157,7 +158,7 @@ public SessionClient build() {
               proxyFactory,
               context);
         } else {
-          proxy = proxyFactory.get();
+          proxy = proxyFactory.get().join();
         }
 
         // If max retries is set, wrap the client in a retrying proxy client.