@@ -104,27 +104,24 @@ public void getWriteConnectionSource(final SingleResultCallback<AsyncConnectionS
104
104
105
105
private void getConnectionSource (final AsyncSupplier <AsyncConnectionSource > connectionSourceSupplier ,
106
106
final SingleResultCallback <AsyncConnectionSource > callback ) {
107
- // wrapper applied at end
108
107
beginAsync ().<AsyncConnectionSource >thenSupply (c -> {
109
108
if (!session .hasActiveTransaction ()) {
110
109
connectionSourceSupplier .getAsync (c );
111
- return ;
112
- }
113
- if (TransactionContext .get (session ) != null ) {
110
+ } else if (TransactionContext .get (session ) != null ) {
114
111
wrapped .getConnectionSource (assertNotNull (session .getPinnedServerAddress ()), c );
115
- return ;
112
+ } else {
113
+ beginAsync ().<AsyncConnectionSource >thenSupply (c2 -> {
114
+ connectionSourceSupplier .getAsync (c2 );
115
+ }).<AsyncConnectionSource >thenApply ((source , c2 ) -> {
116
+ ClusterType clusterType = assertNotNull (source ).getServerDescription ().getClusterType ();
117
+ if (clusterType == SHARDED || clusterType == LOAD_BALANCED ) {
118
+ TransactionContext <AsyncConnection > transactionContext = new TransactionContext <>(clusterType );
119
+ session .setTransactionContext (source .getServerDescription ().getAddress (), transactionContext );
120
+ transactionContext .release (); // The session is responsible for retaining a reference to the context
121
+ }
122
+ c2 .complete (source );
123
+ }).finish (c );
116
124
}
117
- beginAsync ().<AsyncConnectionSource >thenSupply (c2 -> {
118
- connectionSourceSupplier .getAsync (c2 );
119
- }).<AsyncConnectionSource >thenApply ((source , c2 ) -> {
120
- ClusterType clusterType = assertNotNull (source ).getServerDescription ().getClusterType ();
121
- if (clusterType == SHARDED || clusterType == LOAD_BALANCED ) {
122
- TransactionContext <AsyncConnection > transactionContext = new TransactionContext <>(clusterType );
123
- session .setTransactionContext (source .getServerDescription ().getAddress (), transactionContext );
124
- transactionContext .release (); // The session is responsible for retaining a reference to the context
125
- } //
126
- c2 .complete (source );
127
- }).finish (c );
128
125
}).<AsyncConnectionSource >thenApply ((source , c ) -> {
129
126
c .complete (new SessionBindingAsyncConnectionSource (source ));
130
127
}).finish (callback );
@@ -194,12 +191,12 @@ public void getConnection(final SingleResultCallback<AsyncConnection> callback)
194
191
if (transactionContext == null || !transactionContext .isConnectionPinningRequired ()) {
195
192
wrapped .getConnection (c );
196
193
return ;
197
- } //
194
+ }
198
195
AsyncConnection pinnedAsyncConnection = transactionContext .getPinnedConnection ();
199
196
if (pinnedAsyncConnection != null ) {
200
197
c .complete (pinnedAsyncConnection .retain ());
201
198
return ;
202
- } //
199
+ }
203
200
beginAsync ().<AsyncConnection >thenSupply (c2 -> {
204
201
wrapped .getConnection (c2 );
205
202
}).<AsyncConnection >thenApply ((connection , c2 ) -> {
0 commit comments