| 
21 | 21 | import static org.neo4j.driver.internal.reactive.RxUtils.createEmptyPublisher;  | 
22 | 22 | import static org.neo4j.driver.internal.reactive.RxUtils.createSingleItemPublisher;  | 
23 | 23 | 
 
  | 
24 |  | -import java.util.Optional;  | 
25 | 24 | import java.util.Set;  | 
26 | 25 | import java.util.concurrent.CompletableFuture;  | 
27 | 26 | import java.util.concurrent.CompletionException;  | 
@@ -151,65 +150,27 @@ public Set<Bookmark> lastBookmarks() {  | 
151 | 150 |     protected <T> Publisher<T> run(Query query, TransactionConfig config, Function<RxResultCursor, T> cursorToResult) {  | 
152 | 151 |         var cursorPublishFuture = new CompletableFuture<RxResultCursor>();  | 
153 | 152 |         var cursorReference = new AtomicReference<RxResultCursor>();  | 
154 |  | -        return Mono.<T>create(sink -> {  | 
155 |  | -                    var state = new SinkState<T>();  | 
156 |  | -                    sink.onRequest(ignored -> {  | 
157 |  | -                        CompletionStage<T> stage;  | 
158 |  | -                        synchronized (state) {  | 
159 |  | -                            if (state.isCancelled()) {  | 
160 |  | -                                return;  | 
161 |  | -                            }  | 
162 |  | -                            if (state.getStage() != null) {  | 
163 |  | -                                return;  | 
164 |  | -                            }  | 
165 |  | -                            stage = runAsStage(query, config, cursorPublishFuture)  | 
166 |  | -                                    .thenApply(cursor -> {  | 
167 |  | -                                        cursorReference.set(cursor);  | 
168 |  | -                                        return cursor;  | 
169 |  | -                                    })  | 
170 |  | -                                    .thenApply(cursorToResult);  | 
171 |  | -                            state.setStage(stage);  | 
172 |  | -                        }  | 
173 |  | -                        stage.whenComplete((item, completionError) -> {  | 
174 |  | -                            if (completionError == null) {  | 
175 |  | -                                if (item != null) {  | 
176 |  | -                                    sink.success(item);  | 
177 |  | -                                } else {  | 
178 |  | -                                    sink.error(  | 
179 |  | -                                            new IllegalStateException(  | 
180 |  | -                                                    "Unexpected condition, run call has completed successfully with result being null"));  | 
181 |  | -                                }  | 
182 |  | -                            } else {  | 
183 |  | -                                var error = Optional.ofNullable(Futures.completionExceptionCause(completionError))  | 
184 |  | -                                        .orElse(completionError);  | 
185 |  | -                                sink.error(error);  | 
186 |  | -                            }  | 
187 |  | -                        });  | 
188 |  | -                    });  | 
189 |  | -                    sink.onCancel(() -> {  | 
190 |  | -                        CompletionStage<T> stage;  | 
191 |  | -                        synchronized (state) {  | 
192 |  | -                            if (state.isCancelled()) {  | 
193 |  | -                                return;  | 
 | 153 | + | 
 | 154 | +        return createSingleItemPublisher(  | 
 | 155 | +                        () -> runAsStage(query, config, cursorPublishFuture)  | 
 | 156 | +                                .thenApply(cursor -> {  | 
 | 157 | +                                    cursorReference.set(cursor);  | 
 | 158 | +                                    return cursor;  | 
 | 159 | +                                })  | 
 | 160 | +                                .thenApply(cursorToResult),  | 
 | 161 | +                        () -> new IllegalStateException(  | 
 | 162 | +                                "Unexpected condition, run call has completed successfully with result being null"),  | 
 | 163 | +                        value -> {  | 
 | 164 | +                            if (value != null) {  | 
 | 165 | +                                cursorReference.get().rollback().whenComplete((unused, throwable) -> {  | 
 | 166 | +                                    if (throwable != null) {  | 
 | 167 | +                                        cursorPublishFuture.completeExceptionally(throwable);  | 
 | 168 | +                                    } else {  | 
 | 169 | +                                        cursorPublishFuture.complete(null);  | 
 | 170 | +                                    }  | 
 | 171 | +                                });  | 
194 | 172 |                             }  | 
195 |  | -                            state.setCancelled(true);  | 
196 |  | -                            stage = state.getStage();  | 
197 |  | -                        }  | 
198 |  | -                        if (stage != null) {  | 
199 |  | -                            stage.whenComplete((value, ignored) -> {  | 
200 |  | -                                if (value != null) {  | 
201 |  | -                                    cursorReference.get().rollback().whenComplete((unused, throwable) -> {  | 
202 |  | -                                        if (throwable != null) {  | 
203 |  | -                                            cursorPublishFuture.completeExceptionally(throwable);  | 
204 |  | -                                        } else {  | 
205 |  | -                                            cursorPublishFuture.complete(null);  | 
206 |  | -                                        }  | 
207 |  | -                                    });  | 
208 |  | -                                }  | 
209 |  | -                            });  | 
210 |  | -                        }  | 
211 |  | -                    });  | 
212 |  | -                })  | 
 | 173 | +                        })  | 
213 | 174 |                 .doOnNext(value -> cursorPublishFuture.complete(cursorReference.get()))  | 
214 | 175 |                 .doOnError(cursorPublishFuture::completeExceptionally);  | 
215 | 176 |     }  | 
@@ -256,25 +217,4 @@ private <T> CompletionStage<T> releaseConnectionAndRethrow(Throwable throwable)  | 
256 | 217 |     protected <T> Publisher<T> doClose() {  | 
257 | 218 |         return createEmptyPublisher(session::closeAsync);  | 
258 | 219 |     }  | 
259 |  | - | 
260 |  | -    private static class SinkState<T> {  | 
261 |  | -        private CompletionStage<T> stage;  | 
262 |  | -        private boolean cancelled;  | 
263 |  | - | 
264 |  | -        public CompletionStage<T> getStage() {  | 
265 |  | -            return stage;  | 
266 |  | -        }  | 
267 |  | - | 
268 |  | -        public void setStage(CompletionStage<T> stage) {  | 
269 |  | -            this.stage = stage;  | 
270 |  | -        }  | 
271 |  | - | 
272 |  | -        public boolean isCancelled() {  | 
273 |  | -            return cancelled;  | 
274 |  | -        }  | 
275 |  | - | 
276 |  | -        public void setCancelled(boolean cancelled) {  | 
277 |  | -            this.cancelled = cancelled;  | 
278 |  | -        }  | 
279 |  | -    }  | 
280 | 220 | }  | 
0 commit comments