Skip to content

Commit 642b909

Browse files
injectivesinjectives
committed
Add result consumption and disposal to reactive testkit backend session close
1 parent 59bc2b4 commit 642b909

File tree

4 files changed

+123
-7
lines changed

4 files changed

+123
-7
lines changed

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

+2
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

+12
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

+109-1
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,117 @@ 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 ->
62+
sessionHolder.getResultHolder()
63+
.map( this::consumeRequestedDemandAndCancel )
64+
.orElse( Mono.empty() )
65+
.then( Mono.fromDirect( sessionHolder.getSession().close() ) ) )
5666
.then( Mono.just( createResponse() ) );
5767
}
5868

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

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

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

-6
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 );
@@ -79,8 +75,6 @@ public class StartTest implements TestkitRequest
7975
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestRoutingConnectionRecvTimeout\\.test_timeout_unmanaged_tx$", skipMessage );
8076
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestDisconnects\\.test_disconnect_session_on_tx_commit$", skipMessage );
8177
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestSessionRunParameters\\..*$", skipMessage );
82-
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestTxRun\\.test_rollback_tx_on_session_close_unfinished_result$", skipMessage );
83-
REACTIVE_SKIP_PATTERN_TO_REASON.put( "^.*\\.TestTxRun\\.test_rollback_tx_on_session_close_untouched_result$", skipMessage );
8478
}
8579

8680
private StartTestBody data;

0 commit comments

Comments
 (0)