Skip to content

Commit ed8f8b7

Browse files
vbabaninstIncMale
andauthored
Add improved Bulk Write API for Java Reactive Driver (#1583)
- Created and documented the new Reactive Bulk Write API. - Enabled unified and prose tests for reactive Bulk Write API. JAVA-5530 --------- Co-authored-by: Valentin Kovalenko <valentin.male.kovalenko@gmail.com>
1 parent 65e72d0 commit ed8f8b7

File tree

31 files changed

+914
-239
lines changed

31 files changed

+914
-239
lines changed

driver-core/src/main/com/mongodb/internal/async/AsyncBatchCursor.java

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,11 @@
1919
import com.mongodb.internal.operation.BatchCursor;
2020

2121
import java.io.Closeable;
22+
import java.util.ArrayList;
2223
import java.util.List;
2324

25+
import static com.mongodb.internal.async.AsyncRunnable.beginAsync;
26+
2427
/**
2528
* MongoDB returns query results as batches, and this interface provides an asynchronous iterator over those batches. The first call to
2629
* the {@code next} method will return the first batch, and subsequent calls will trigger an asynchronous request to get the next batch
@@ -72,4 +75,22 @@ public interface AsyncBatchCursor<T> extends Closeable {
7275
*/
7376
@Override
7477
void close();
78+
79+
default void exhaust(final SingleResultCallback<List<List<T>>> finalCallback) {
80+
List<List<T>> results = new ArrayList<>();
81+
82+
beginAsync().thenRunDoWhileLoop(iterationCallback -> {
83+
beginAsync().<List<T>>thenSupply(c -> {
84+
next(c);
85+
}).thenConsume((batch, c) -> {
86+
if (!batch.isEmpty()) {
87+
results.add(batch);
88+
}
89+
c.complete(c);
90+
}).finish(iterationCallback);
91+
}, () -> !this.isClosed()
92+
).<List<List<T>>>thenSupply(c -> {
93+
c.complete(results);
94+
}).finish(finalCallback);
95+
}
7596
}

driver-core/src/main/com/mongodb/internal/async/AsyncRunnable.java

Lines changed: 32 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,12 @@
1717
package com.mongodb.internal.async;
1818

1919
import com.mongodb.internal.TimeoutContext;
20+
import com.mongodb.internal.async.function.AsyncCallbackLoop;
21+
import com.mongodb.internal.async.function.LoopState;
2022
import com.mongodb.internal.async.function.RetryState;
2123
import com.mongodb.internal.async.function.RetryingAsyncCallbackSupplier;
2224

25+
import java.util.function.BooleanSupplier;
2326
import java.util.function.Predicate;
2427
import java.util.function.Supplier;
2528

@@ -120,49 +123,6 @@ static AsyncRunnable beginAsync() {
120123
return (c) -> c.complete(c);
121124
}
122125

123-
/**
124-
* Must be invoked at end of async chain
125-
* @param runnable the sync code to invoke (under non-exceptional flow)
126-
* prior to the callback
127-
* @param callback the callback provided by the method the chain is used in
128-
*/
129-
default void thenRunAndFinish(final Runnable runnable, final SingleResultCallback<Void> callback) {
130-
this.finish((r, e) -> {
131-
if (e != null) {
132-
callback.completeExceptionally(e);
133-
return;
134-
}
135-
try {
136-
runnable.run();
137-
} catch (Throwable t) {
138-
callback.completeExceptionally(t);
139-
return;
140-
}
141-
callback.complete(callback);
142-
});
143-
}
144-
145-
/**
146-
* See {@link #thenRunAndFinish(Runnable, SingleResultCallback)}, but the runnable
147-
* will always be executed, including on the exceptional path.
148-
* @param runnable the runnable
149-
* @param callback the callback
150-
*/
151-
default void thenAlwaysRunAndFinish(final Runnable runnable, final SingleResultCallback<Void> callback) {
152-
this.finish((r, e) -> {
153-
try {
154-
runnable.run();
155-
} catch (Throwable t) {
156-
if (e != null) {
157-
t.addSuppressed(e);
158-
}
159-
callback.completeExceptionally(t);
160-
return;
161-
}
162-
callback.onResult(r, e);
163-
});
164-
}
165-
166126
/**
167127
* @param runnable The async runnable to run after this runnable
168128
* @return the composition of this runnable and the runnable, a runnable
@@ -282,4 +242,33 @@ default AsyncRunnable thenRunRetryingWhile(
282242
).get(callback);
283243
});
284244
}
245+
246+
/**
247+
* This method is equivalent to a do-while loop, where the loop body is executed first and
248+
* then the condition is checked to determine whether the loop should continue.
249+
*
250+
* @param loopBodyRunnable the asynchronous task to be executed in each iteration of the loop
251+
* @param whileCheck a condition to check after each iteration; the loop continues as long as this condition returns true
252+
* @return the composition of this and the looping branch
253+
* @see AsyncCallbackLoop
254+
*/
255+
default AsyncRunnable thenRunDoWhileLoop(final AsyncRunnable loopBodyRunnable, final BooleanSupplier whileCheck) {
256+
return thenRun(finalCallback -> {
257+
LoopState loopState = new LoopState();
258+
new AsyncCallbackLoop(loopState, iterationCallback -> {
259+
260+
loopBodyRunnable.finish((result, t) -> {
261+
if (t != null) {
262+
iterationCallback.completeExceptionally(t);
263+
return;
264+
}
265+
if (loopState.breakAndCompleteIf(() -> !whileCheck.getAsBoolean(), iterationCallback)) {
266+
return;
267+
}
268+
iterationCallback.complete(iterationCallback);
269+
});
270+
271+
}).run(finalCallback);
272+
});
273+
}
285274
}

driver-core/src/main/com/mongodb/internal/async/AsyncSupplier.java

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,49 @@ default void finish(final SingleResultCallback<T> callback) {
8181
}
8282
}
8383

84+
/**
85+
* Must be invoked at end of async chain
86+
* @param runnable the sync code to invoke (under non-exceptional flow)
87+
* prior to the callback
88+
* @param callback the callback provided by the method the chain is used in
89+
*/
90+
default void thenRunAndFinish(final Runnable runnable, final SingleResultCallback<T> callback) {
91+
this.finish((r, e) -> {
92+
if (e != null) {
93+
callback.completeExceptionally(e);
94+
return;
95+
}
96+
try {
97+
runnable.run();
98+
} catch (Throwable t) {
99+
callback.completeExceptionally(t);
100+
return;
101+
}
102+
callback.onResult(r, null);
103+
});
104+
}
105+
106+
/**
107+
* See {@link #thenRunAndFinish(Runnable, SingleResultCallback)}, but the runnable
108+
* will always be executed, including on the exceptional path.
109+
* @param runnable the runnable
110+
* @param callback the callback
111+
*/
112+
default void thenAlwaysRunAndFinish(final Runnable runnable, final SingleResultCallback<T> callback) {
113+
this.finish((r, e) -> {
114+
try {
115+
runnable.run();
116+
} catch (Throwable t) {
117+
if (e != null) {
118+
t.addSuppressed(e);
119+
}
120+
callback.completeExceptionally(t);
121+
return;
122+
}
123+
callback.onResult(r, e);
124+
});
125+
}
126+
84127
/**
85128
* @param function The async function to run after this supplier
86129
* @return the composition of this supplier and the function, a supplier
Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
/*
2+
* Copyright 2008-present MongoDB, Inc.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.mongodb.internal.async;
17+
18+
import com.mongodb.annotations.NotThreadSafe;
19+
import com.mongodb.lang.Nullable;
20+
21+
import static com.mongodb.assertions.Assertions.assertNotNull;
22+
23+
@NotThreadSafe
24+
public final class MutableValue<T> {
25+
private T value;
26+
27+
public MutableValue(@Nullable final T value) {
28+
this.value = value;
29+
}
30+
31+
public MutableValue() {
32+
this(null);
33+
}
34+
35+
public T get() {
36+
return assertNotNull(value);
37+
}
38+
39+
@Nullable
40+
public T getNullable() {
41+
return value;
42+
}
43+
44+
public void set(@Nullable final T value) {
45+
this.value = value;
46+
}
47+
}

driver-core/src/main/com/mongodb/internal/async/function/AsyncCallbackSupplier.java

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
*/
1616
package com.mongodb.internal.async.function;
1717

18-
import com.mongodb.annotations.NotThreadSafe;
18+
import com.mongodb.internal.async.MutableValue;
1919
import com.mongodb.internal.async.SingleResultCallback;
2020

2121
import java.util.function.Supplier;
@@ -68,16 +68,12 @@ public interface AsyncCallbackSupplier<R> {
6868
* This is a price we have to pay to provide a guarantee similar to that of the {@code finally} block.
6969
*/
7070
default AsyncCallbackSupplier<R> whenComplete(final Runnable after) {
71-
@NotThreadSafe
72-
final class MutableBoolean {
73-
private boolean value;
74-
}
75-
MutableBoolean afterExecuted = new MutableBoolean();
71+
MutableValue<Boolean> afterExecuted = new MutableValue<>(false);
7672
Runnable trackableAfter = () -> {
7773
try {
7874
after.run();
7975
} finally {
80-
afterExecuted.value = true;
76+
afterExecuted.set(true);
8177
}
8278
};
8379
return callback -> {
@@ -103,7 +99,7 @@ final class MutableBoolean {
10399
primaryUnexpectedException = unexpectedException;
104100
throw unexpectedException;
105101
} finally {
106-
if (primaryUnexpectedException != null && !afterExecuted.value) {
102+
if (primaryUnexpectedException != null && !afterExecuted.get()) {
107103
try {
108104
trackableAfter.run();
109105
} catch (Throwable afterException) {

driver-core/src/main/com/mongodb/internal/operation/AsyncOperationHelper.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -344,7 +344,7 @@ static <T> CommandReadTransformerAsync<BsonDocument, AsyncBatchCursor<T>> asyncS
344344
}
345345

346346
static <T> AsyncBatchCursor<T> cursorDocumentToAsyncBatchCursor(final TimeoutMode timeoutMode, final BsonDocument cursorDocument,
347-
final int batchSize, final Decoder<T> decoder, final BsonValue comment, final AsyncConnectionSource source,
347+
final int batchSize, final Decoder<T> decoder, @Nullable final BsonValue comment, final AsyncConnectionSource source,
348348
final AsyncConnection connection) {
349349
return new AsyncCommandBatchCursor<>(timeoutMode, cursorDocument, batchSize, 0, decoder, comment, source, connection);
350350
}

driver-core/src/main/com/mongodb/internal/operation/AsyncOperations.java

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,9 @@
4444
import com.mongodb.client.model.SearchIndexModel;
4545
import com.mongodb.client.model.UpdateOptions;
4646
import com.mongodb.client.model.WriteModel;
47+
import com.mongodb.client.model.bulk.ClientBulkWriteOptions;
48+
import com.mongodb.client.model.bulk.ClientBulkWriteResult;
49+
import com.mongodb.client.model.bulk.ClientNamespacedWriteModel;
4750
import com.mongodb.client.model.changestream.FullDocument;
4851
import com.mongodb.client.model.changestream.FullDocumentBeforeChange;
4952
import com.mongodb.internal.TimeoutSettings;
@@ -293,6 +296,12 @@ public AsyncWriteOperation<BulkWriteResult> bulkWrite(final List<? extends Write
293296
return operations.bulkWrite(requests, options);
294297
}
295298

299+
public AsyncWriteOperation<ClientBulkWriteResult> clientBulkWriteOperation(
300+
final List<? extends ClientNamespacedWriteModel> clientWriteModels,
301+
@Nullable final ClientBulkWriteOptions options) {
302+
return operations.clientBulkWriteOperation(clientWriteModels, options);
303+
}
304+
296305
public <TResult> AsyncReadOperation<TResult> commandRead(final Bson command, final Class<TResult> resultClass) {
297306
return operations.commandRead(command, resultClass);
298307
}

driver-core/src/main/com/mongodb/internal/operation/BatchCursor.java

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,12 @@
2525
import java.util.Iterator;
2626
import java.util.List;
2727

28+
import static java.util.Spliterator.IMMUTABLE;
29+
import static java.util.Spliterator.ORDERED;
30+
import static java.util.Spliterators.spliteratorUnknownSize;
31+
import static java.util.stream.Collectors.toList;
32+
import static java.util.stream.StreamSupport.stream;
33+
2834
/**
2935
* MongoDB returns query results as batches, and this interface provideds an iterator over those batches. The first call to
3036
* the {@code next} method will return the first batch, and subsequent calls will trigger a request to get the next batch
@@ -98,4 +104,9 @@ public interface BatchCursor<T> extends Iterator<List<T>>, Closeable {
98104
ServerCursor getServerCursor();
99105

100106
ServerAddress getServerAddress();
107+
108+
default List<List<T>> exhaust() {
109+
return stream(spliteratorUnknownSize(this, ORDERED | IMMUTABLE), false)
110+
.collect(toList());
111+
}
101112
}

0 commit comments

Comments
 (0)