Skip to content

Commit b673c01

Browse files
injectivesinjectives
authored
Add result consumption and disposal to reactive testkit backend session close (#1010) (#1016)
1 parent e1f99b1 commit b673c01

File tree

4 files changed

+127
-8
lines changed

4 files changed

+127
-8
lines changed

‎testkit-backend/src/main/java/neo4j/org/testkit/backend/holder/RxResultHolder.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,11 +38,13 @@ public class RxResultHolder extends AbstractResultHolder<RxSessionHolder,RxTrans
3838
public RxResultHolder( RxSessionHolder sessionHolder, RxResult result )
3939
{
4040
super( sessionHolder, result );
41+
sessionHolder.setResultHolder( this );
4142
}
4243

4344
public RxResultHolder( RxTransactionHolder transactionHolder, RxResult result )
4445
{
4546
super( transactionHolder, result );
47+
transactionHolder.getSessionHolder().setResultHolder( this );
4648
}
4749

4850
public Optional<RxBlockingSubscriber<Record>> getSubscriber()

‎testkit-backend/src/main/java/neo4j/org/testkit/backend/holder/RxSessionHolder.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,25 @@
1818
*/
1919
package neo4j.org.testkit.backend.holder;
2020

21+
import lombok.Setter;
22+
23+
import java.util.Optional;
24+
2125
import org.neo4j.driver.SessionConfig;
2226
import org.neo4j.driver.reactive.RxSession;
2327

2428
public class RxSessionHolder extends AbstractSessionHolder<RxSession>
2529
{
30+
@Setter
31+
private RxResultHolder resultHolder;
32+
2633
public RxSessionHolder( DriverHolder driverHolder, RxSession session, SessionConfig config )
2734
{
2835
super( driverHolder, session, config );
2936
}
37+
38+
public Optional<RxResultHolder> getResultHolder()
39+
{
40+
return Optional.ofNullable( resultHolder );
41+
}
3042
}

‎testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/SessionClose.java‎

Lines changed: 113 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,12 +20,18 @@
2020

2121
import lombok.Getter;
2222
import lombok.Setter;
23+
import neo4j.org.testkit.backend.RxBlockingSubscriber;
2324
import neo4j.org.testkit.backend.TestkitState;
25+
import neo4j.org.testkit.backend.holder.RxResultHolder;
2426
import neo4j.org.testkit.backend.messages.responses.Session;
2527
import neo4j.org.testkit.backend.messages.responses.TestkitResponse;
2628
import reactor.core.publisher.Mono;
2729

30+
import java.util.concurrent.CompletableFuture;
2831
import java.util.concurrent.CompletionStage;
32+
import java.util.concurrent.atomic.AtomicLong;
33+
34+
import org.neo4j.driver.Record;
2935

3036
@Setter
3137
@Getter
@@ -52,15 +58,121 @@ public CompletionStage<TestkitResponse> processAsync( TestkitState testkitState
5258
public Mono<TestkitResponse> processRx( TestkitState testkitState )
5359
{
5460
return testkitState.getRxSessionHolder( data.getSessionId() )
55-
.flatMap( sessionHolder -> Mono.fromDirect( sessionHolder.getSession().close() ) )
61+
.flatMap( sessionHolder -> sessionHolder.getResultHolder()
62+
.map( this::consumeRequestedDemandAndCancelIfSubscribed )
63+
.orElse( Mono.empty() )
64+
.then( Mono.fromDirect( sessionHolder.getSession().close() ) ) )
5665
.then( Mono.just( createResponse() ) );
5766
}
5867

68+
private Mono<Void> consumeRequestedDemandAndCancelIfSubscribed( RxResultHolder resultHolder )
69+
{
70+
return resultHolder.getSubscriber()
71+
.map( subscriber -> Mono.fromCompletionStage( consumeRequestedDemandAndCancelIfSubscribed( resultHolder, subscriber ) ) )
72+
.orElse( Mono.empty() );
73+
}
74+
75+
private CompletionStage<Void> consumeRequestedDemandAndCancelIfSubscribed( RxResultHolder resultHolder, RxBlockingSubscriber<Record> subscriber )
76+
{
77+
if ( subscriber.getCompletionStage().toCompletableFuture().isDone() )
78+
{
79+
return CompletableFuture.completedFuture( null );
80+
}
81+
82+
return new DemandConsumer<>( subscriber, resultHolder.getRequestedRecordsCounter() )
83+
.getCompletedStage()
84+
.thenCompose( completionReason ->
85+
{
86+
CompletionStage<Void> result;
87+
switch ( completionReason )
88+
{
89+
case REQUESTED_DEMAND_CONSUMED:
90+
result = subscriber.getSubscriptionStage().thenApply( subscription ->
91+
{
92+
subscription.cancel();
93+
return null;
94+
} );
95+
break;
96+
case RECORD_STREAM_EXHAUSTED:
97+
result = CompletableFuture.completedFuture( null );
98+
break;
99+
default:
100+
result = new CompletableFuture<>();
101+
result.toCompletableFuture()
102+
.completeExceptionally( new RuntimeException( "Unexpected completion reason: " + completionReason ) );
103+
}
104+
return result;
105+
} );
106+
}
107+
59108
private Session createResponse()
60109
{
61110
return Session.builder().data( Session.SessionBody.builder().id( data.getSessionId() ).build() ).build();
62111
}
63112

113+
private static class DemandConsumer<T>
114+
{
115+
private final RxBlockingSubscriber<T> subscriber;
116+
private final AtomicLong unfulfilledDemandCounter;
117+
@Getter
118+
private final CompletableFuture<CompletionReason> completedStage = new CompletableFuture<>();
119+
120+
private enum CompletionReason
121+
{
122+
REQUESTED_DEMAND_CONSUMED,
123+
RECORD_STREAM_EXHAUSTED
124+
}
125+
126+
private DemandConsumer( RxBlockingSubscriber<T> subscriber, AtomicLong unfulfilledDemandCounter )
127+
{
128+
this.subscriber = subscriber;
129+
this.unfulfilledDemandCounter = unfulfilledDemandCounter;
130+
131+
subscriber.getCompletionStage().whenComplete( this::onComplete );
132+
if ( this.unfulfilledDemandCounter.get() > 0 )
133+
{
134+
setupNextSignalConsumer();
135+
}
136+
}
137+
138+
private void setupNextSignalConsumer()
139+
{
140+
CompletableFuture<T> consumer = new CompletableFuture<>();
141+
subscriber.setNextSignalConsumer( consumer );
142+
consumer.whenComplete( this::onNext );
143+
}
144+
145+
private void onNext( T ignored, Throwable throwable )
146+
{
147+
if ( throwable != null )
148+
{
149+
completedStage.completeExceptionally( throwable );
150+
return;
151+
}
152+
153+
if ( unfulfilledDemandCounter.decrementAndGet() > 0 )
154+
{
155+
setupNextSignalConsumer();
156+
}
157+
else
158+
{
159+
completedStage.complete( CompletionReason.REQUESTED_DEMAND_CONSUMED );
160+
}
161+
}
162+
163+
private void onComplete( Void ignored, Throwable throwable )
164+
{
165+
if ( throwable != null )
166+
{
167+
completedStage.completeExceptionally( throwable );
168+
}
169+
else
170+
{
171+
completedStage.complete( CompletionReason.RECORD_STREAM_EXHAUSTED );
172+
}
173+
}
174+
}
175+
64176
@Setter
65177
@Getter
66178
private static class SessionCloseBody

‎testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/StartTest.java‎

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -61,10 +61,6 @@ public class StartTest implements TestkitRequest
6161
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestDisconnects\\.test_disconnect_after_hello$", skipMessage );
6262
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestDisconnects\\.test_disconnect_session_on_run$", skipMessage );
6363
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestDisconnects\\.test_disconnect_on_tx_run$", skipMessage );
64-
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestRetry\\..*$", "Unfinished results consumption" );
65-
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestRetryClustering\\..*$", "Unfinished results consumption" );
66-
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestSessionRun\\.test_discard_on_session_close_unfinished_result$",
67-
"Does not support partially consumed state" );
6864
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.NoRouting[^.]+\\.test_should_error_on_database_shutdown_using_tx_run$", "Session close throws error" );
6965
skipMessage = "Requires investigation";
7066
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestProtocolVersions\\.test_server_agent", skipMessage );
@@ -81,12 +77,9 @@ public class StartTest implements TestkitRequest
8177
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestDisconnects\\.test_disconnect_session_on_tx_commit$", skipMessage );
8278
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestDisconnects\\.test_fail_on_reset$", skipMessage );
8379
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestSessionRunParameters\\..*$", skipMessage );
84-
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestTxRun\\.test_rollback_tx_on_session_close_unfinished_result$", skipMessage );
85-
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestTxRun\\.test_rollback_tx_on_session_close_untouched_result$", skipMessage );
8680
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestSessionRun\\.test_autocommit_transactions_should_support_timeout$", skipMessage );
8781
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestSessionRun\\.test_fails_on_bad_syntax$", skipMessage );
8882
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestSessionRun\\.test_fails_on_missing_parameter$", skipMessage );
89-
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestSessionRun\\.test_iteration_larger_than_fetch_size$", skipMessage );
9083
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestSessionRun\\.test_iteration_nested$", skipMessage );
9184
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestTxFuncRun\\.test_iteration_nested$", skipMessage );
9285
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestTxRun\\.test_should_not_run_valid_query_in_invalid_tx$", skipMessage );

0 commit comments

Comments
 (0)