18
18
*/
19
19
package org .neo4j .driver .internal ;
20
20
21
+ import java .util .ArrayList ;
22
+ import java .util .List ;
21
23
import java .util .Map ;
22
24
import java .util .concurrent .CompletableFuture ;
25
+ import java .util .concurrent .CompletionException ;
23
26
import java .util .concurrent .CompletionStage ;
24
27
import java .util .function .BiConsumer ;
28
+ import java .util .function .BiFunction ;
25
29
30
+ import org .neo4j .driver .internal .async .InternalStatementResultCursor ;
26
31
import org .neo4j .driver .internal .async .QueryRunner ;
27
32
import org .neo4j .driver .internal .handlers .BeginTxResponseHandler ;
28
33
import org .neo4j .driver .internal .handlers .CommitTxResponseHandler ;
29
34
import org .neo4j .driver .internal .handlers .NoOpResponseHandler ;
30
35
import org .neo4j .driver .internal .handlers .RollbackTxResponseHandler ;
31
36
import org .neo4j .driver .internal .spi .Connection ;
37
+ import org .neo4j .driver .internal .spi .ResponseHandler ;
32
38
import org .neo4j .driver .internal .types .InternalTypeSystem ;
39
+ import org .neo4j .driver .internal .util .Futures ;
33
40
import org .neo4j .driver .v1 .Record ;
34
41
import org .neo4j .driver .v1 .Session ;
35
42
import org .neo4j .driver .v1 .Statement ;
43
50
44
51
import static java .util .Collections .emptyMap ;
45
52
import static java .util .concurrent .CompletableFuture .completedFuture ;
53
+ import static org .neo4j .driver .internal .util .Futures .completionErrorCause ;
46
54
import static org .neo4j .driver .internal .util .Futures .failedFuture ;
47
55
import static org .neo4j .driver .internal .util .Futures .getBlocking ;
48
56
import static org .neo4j .driver .v1 .Values .value ;
@@ -86,6 +94,7 @@ private enum State
86
94
private final Connection connection ;
87
95
private final NetworkSession session ;
88
96
97
+ private final List <CompletionStage <InternalStatementResultCursor >> resultCursors = new ArrayList <>();
89
98
private volatile Bookmark bookmark = Bookmark .empty ();
90
99
private volatile State state = State .ACTIVE ;
91
100
@@ -169,7 +178,9 @@ else if ( state == State.TERMINATED )
169
178
}
170
179
else
171
180
{
172
- return doCommitAsync ().whenComplete ( transactionClosed ( State .COMMITTED ) );
181
+ return receiveFailures ()
182
+ .thenCompose ( failure -> doCommitAsync ().handle ( handleCommitOrRollback ( failure ) ) )
183
+ .whenComplete ( transactionClosed ( State .COMMITTED ) );
173
184
}
174
185
}
175
186
@@ -192,38 +203,12 @@ else if ( state == State.TERMINATED )
192
203
}
193
204
else
194
205
{
195
- return doRollbackAsync ().whenComplete ( transactionClosed ( State .ROLLED_BACK ) );
206
+ return receiveFailures ()
207
+ .thenCompose ( failure -> doRollbackAsync ().handle ( handleCommitOrRollback ( failure ) ) )
208
+ .whenComplete ( transactionClosed ( State .ROLLED_BACK ) );
196
209
}
197
210
}
198
211
199
- private BiConsumer <Void ,Throwable > transactionClosed ( State newState )
200
- {
201
- return ( ignore , error ) ->
202
- {
203
- state = newState ;
204
- connection .releaseInBackground ();
205
- session .setBookmark ( bookmark );
206
- };
207
- }
208
-
209
- private CompletionStage <Void > doCommitAsync ()
210
- {
211
- CompletableFuture <Void > commitFuture = new CompletableFuture <>();
212
- connection .runAndFlush ( COMMIT_QUERY , emptyMap (), NoOpResponseHandler .INSTANCE ,
213
- new CommitTxResponseHandler ( commitFuture , this ) );
214
-
215
- return commitFuture .thenRun ( () -> state = State .COMMITTED );
216
- }
217
-
218
- private CompletionStage <Void > doRollbackAsync ()
219
- {
220
- CompletableFuture <Void > rollbackFuture = new CompletableFuture <>();
221
- connection .runAndFlush ( ROLLBACK_QUERY , emptyMap (), NoOpResponseHandler .INSTANCE ,
222
- new RollbackTxResponseHandler ( rollbackFuture ) );
223
-
224
- return rollbackFuture .thenRun ( () -> state = State .ROLLED_BACK );
225
- }
226
-
227
212
@ Override
228
213
public StatementResult run ( String statementText , Value statementParameters )
229
214
{
@@ -280,23 +265,31 @@ public CompletionStage<StatementResultCursor> runAsync( String statementTemplate
280
265
@ Override
281
266
public StatementResult run ( Statement statement )
282
267
{
283
- ensureCanRunQueries ();
284
- StatementResultCursor cursor = getBlocking ( QueryRunner .runAsBlocking ( connection , statement , this ) );
268
+ StatementResultCursor cursor = getBlocking ( run ( statement , false ) );
285
269
return new InternalStatementResult ( cursor );
286
270
}
287
271
288
272
@ Override
289
273
public CompletionStage <StatementResultCursor > runAsync ( Statement statement )
290
274
{
291
- ensureCanRunQueries ();
292
275
//noinspection unchecked
293
- return (CompletionStage ) QueryRunner . runAsAsync ( connection , statement , this );
276
+ return (CompletionStage ) run ( statement , true );
294
277
}
295
278
296
- @ Override
297
- public boolean isOpen ()
279
+ private CompletionStage <InternalStatementResultCursor > run ( Statement statement , boolean asAsync )
298
280
{
299
- return state .txOpen ;
281
+ ensureCanRunQueries ();
282
+ CompletionStage <InternalStatementResultCursor > result ;
283
+ if ( asAsync )
284
+ {
285
+ result = QueryRunner .runAsAsync ( connection , statement , this );
286
+ }
287
+ else
288
+ {
289
+ result = QueryRunner .runAsBlocking ( connection , statement , this );
290
+ }
291
+ resultCursors .add ( result );
292
+ return result ;
300
293
}
301
294
302
295
private void ensureCanRunQueries ()
@@ -324,6 +317,12 @@ else if ( state == State.TERMINATED )
324
317
}
325
318
}
326
319
320
+ @ Override
321
+ public boolean isOpen ()
322
+ {
323
+ return state .txOpen ;
324
+ }
325
+
327
326
@ Override
328
327
public TypeSystem typeSystem ()
329
328
{
@@ -347,4 +346,56 @@ public void setBookmark( Bookmark bookmark )
347
346
this .bookmark = bookmark ;
348
347
}
349
348
}
349
+
350
+ private CompletionStage <Void > doCommitAsync ()
351
+ {
352
+ CompletableFuture <Void > commitFuture = new CompletableFuture <>();
353
+ ResponseHandler pullAllHandler = new CommitTxResponseHandler ( commitFuture , this );
354
+ connection .runAndFlush ( COMMIT_QUERY , emptyMap (), NoOpResponseHandler .INSTANCE , pullAllHandler );
355
+ return commitFuture ;
356
+ }
357
+
358
+ private CompletionStage <Void > doRollbackAsync ()
359
+ {
360
+ CompletableFuture <Void > rollbackFuture = new CompletableFuture <>();
361
+ ResponseHandler pullAllHandler = new RollbackTxResponseHandler ( rollbackFuture );
362
+ connection .runAndFlush ( ROLLBACK_QUERY , emptyMap (), NoOpResponseHandler .INSTANCE , pullAllHandler );
363
+ return rollbackFuture ;
364
+ }
365
+
366
+ private BiFunction <Void ,Throwable ,Void > handleCommitOrRollback ( Throwable cursorFailure )
367
+ {
368
+ return ( ignore , commitOrRollbackError ) ->
369
+ {
370
+ if ( cursorFailure != null )
371
+ {
372
+ throw new CompletionException ( completionErrorCause ( cursorFailure ) );
373
+ }
374
+ else if ( commitOrRollbackError != null )
375
+ {
376
+ throw new CompletionException ( completionErrorCause ( commitOrRollbackError ) );
377
+ }
378
+ else
379
+ {
380
+ return null ;
381
+ }
382
+ };
383
+ }
384
+
385
+ private BiConsumer <Object ,Throwable > transactionClosed ( State newState )
386
+ {
387
+ return ( ignore , error ) ->
388
+ {
389
+ state = newState ;
390
+ connection .releaseInBackground ();
391
+ session .setBookmark ( bookmark );
392
+ };
393
+ }
394
+
395
+ private CompletionStage <Throwable > receiveFailures ()
396
+ {
397
+ return resultCursors .stream ()
398
+ .map ( stage -> stage .thenCompose ( InternalStatementResultCursor ::failureAsync ) )
399
+ .reduce ( completedFuture ( null ), Futures ::firstNotNull );
400
+ }
350
401
}
0 commit comments