24
24
import com .mongodb .internal .binding .AsyncConnectionSource ;
25
25
import com .mongodb .internal .binding .AsyncReadBinding ;
26
26
import com .mongodb .internal .operation .OperationHelper .AsyncCallableWithSource ;
27
+ import com .mongodb .lang .NonNull ;
27
28
import org .bson .BsonDocument ;
28
29
import org .bson .BsonTimestamp ;
29
30
import org .bson .RawBsonDocument ;
30
31
31
32
import java .util .ArrayList ;
32
33
import java .util .List ;
33
34
import java .util .concurrent .atomic .AtomicBoolean ;
35
+ import java .util .concurrent .atomic .AtomicReference ;
34
36
37
+ import static com .mongodb .assertions .Assertions .assertNotNull ;
38
+ import static com .mongodb .assertions .Assertions .assertNull ;
35
39
import static com .mongodb .internal .async .ErrorHandlingResultCallback .errorHandlingCallback ;
36
40
import static com .mongodb .internal .operation .ChangeStreamBatchCursorHelper .isRetryableError ;
37
41
import static com .mongodb .internal .operation .OperationHelper .LOGGER ;
@@ -44,7 +48,12 @@ final class AsyncChangeStreamBatchCursor<T> implements AsyncAggregateResponseBat
44
48
private final int maxWireVersion ;
45
49
46
50
private volatile BsonDocument resumeToken ;
47
- private volatile AsyncAggregateResponseBatchCursor <RawBsonDocument > wrapped ;
51
+ /**
52
+ * {@linkplain ChangeStreamBatchCursorHelper#isRetryableError(Throwable, int) Retryable errors} can result in
53
+ * {@code wrapped} containing {@code null} and {@link #isClosed} being {@code false}.
54
+ * This represents a situation in which the wrapped object was closed by {@code this} but {@code this} remained open.
55
+ */
56
+ private final AtomicReference <AsyncAggregateResponseBatchCursor <RawBsonDocument >> wrapped ;
48
57
private final AtomicBoolean isClosed ;
49
58
50
59
AsyncChangeStreamBatchCursor (final ChangeStreamOperation <T > changeStreamOperation ,
@@ -53,16 +62,17 @@ final class AsyncChangeStreamBatchCursor<T> implements AsyncAggregateResponseBat
53
62
final BsonDocument resumeToken ,
54
63
final int maxWireVersion ) {
55
64
this .changeStreamOperation = changeStreamOperation ;
56
- this .wrapped = wrapped ;
65
+ this .wrapped = new AtomicReference <>( assertNotNull ( wrapped )) ;
57
66
this .binding = binding ;
58
67
binding .retain ();
59
68
this .resumeToken = resumeToken ;
60
69
this .maxWireVersion = maxWireVersion ;
61
70
isClosed = new AtomicBoolean ();
62
71
}
63
72
73
+ @ NonNull
64
74
AsyncAggregateResponseBatchCursor <RawBsonDocument > getWrapped () {
65
- return wrapped ;
75
+ return assertNotNull ( wrapped . get ()) ;
66
76
}
67
77
68
78
@ Override
@@ -93,7 +103,7 @@ public void apply(final AsyncAggregateResponseBatchCursor<RawBsonDocument> curso
93
103
public void close () {
94
104
if (isClosed .compareAndSet (false , true )) {
95
105
try {
96
- wrapped . close ();
106
+ nullifyAndCloseWrapped ();
97
107
} finally {
98
108
binding .release ();
99
109
}
@@ -102,22 +112,63 @@ public void close() {
102
112
103
113
@ Override
104
114
public void setBatchSize (final int batchSize ) {
105
- wrapped .setBatchSize (batchSize );
115
+ getWrapped () .setBatchSize (batchSize );
106
116
}
107
117
108
118
@ Override
109
119
public int getBatchSize () {
110
- return wrapped .getBatchSize ();
120
+ return getWrapped () .getBatchSize ();
111
121
}
112
122
113
123
@ Override
114
124
public boolean isClosed () {
115
- return isClosed .get ();
125
+ if (isClosed .get ()) {
126
+ return true ;
127
+ } else if (wrappedClosedItself ()) {
128
+ close ();
129
+ return true ;
130
+ } else {
131
+ return false ;
132
+ }
133
+ }
134
+
135
+ private boolean wrappedClosedItself () {
136
+ AsyncAggregateResponseBatchCursor <RawBsonDocument > observedWrapped = wrapped .get ();
137
+ return observedWrapped != null && observedWrapped .isClosed ();
138
+ }
139
+
140
+ /**
141
+ * {@code null} is written to {@link #wrapped} before closing the wrapped object to maintain the following guarantee:
142
+ * if {@link #wrappedClosedItself()} observes a {@linkplain AsyncAggregateResponseBatchCursor#isClosed() closed} wrapped object,
143
+ * then it closed itself as opposed to being closed by {@code this}.
144
+ */
145
+ private void nullifyAndCloseWrapped () {
146
+ AsyncAggregateResponseBatchCursor <RawBsonDocument > observedWrapped = wrapped .getAndSet (null );
147
+ if (observedWrapped != null ) {
148
+ observedWrapped .close ();
149
+ }
150
+ }
151
+
152
+ /**
153
+ * This method guarantees that the {@code newValue} argument is closed even if
154
+ * {@link #setWrappedOrCloseIt(AsyncAggregateResponseBatchCursor)} is called concurrently with or after (in the happens-before order)
155
+ * the method {@link #close()}.
156
+ */
157
+ private void setWrappedOrCloseIt (final AsyncAggregateResponseBatchCursor <RawBsonDocument > newValue ) {
158
+ if (isClosed ()) {
159
+ assertNull (this .wrapped .get ());
160
+ newValue .close ();
161
+ } else {
162
+ assertNull (this .wrapped .getAndSet (newValue ));
163
+ if (isClosed ()) {
164
+ nullifyAndCloseWrapped ();
165
+ }
166
+ }
116
167
}
117
168
118
169
@ Override
119
170
public BsonDocument getPostBatchResumeToken () {
120
- return wrapped .getPostBatchResumeToken ();
171
+ return getWrapped () .getPostBatchResumeToken ();
121
172
}
122
173
123
174
@ Override
@@ -127,7 +178,7 @@ public BsonTimestamp getOperationTime() {
127
178
128
179
@ Override
129
180
public boolean isFirstBatchEmpty () {
130
- return wrapped .isFirstBatchEmpty ();
181
+ return getWrapped () .isFirstBatchEmpty ();
131
182
}
132
183
133
184
@ Override
@@ -178,18 +229,18 @@ private interface AsyncBlock {
178
229
179
230
private void resumeableOperation (final AsyncBlock asyncBlock , final SingleResultCallback <List <RawBsonDocument >> callback ,
180
231
final boolean tryNext ) {
181
- if (isClosed . get ()) {
232
+ if (isClosed ()) {
182
233
callback .onResult (null , new MongoException (format ("%s called after the cursor was closed." ,
183
234
tryNext ? "tryNext()" : "next()" )));
184
235
return ;
185
236
}
186
- asyncBlock .apply (wrapped , new SingleResultCallback <List <RawBsonDocument >>() {
237
+ asyncBlock .apply (getWrapped () , new SingleResultCallback <List <RawBsonDocument >>() {
187
238
@ Override
188
239
public void onResult (final List <RawBsonDocument > result , final Throwable t ) {
189
240
if (t == null ) {
190
241
callback .onResult (result , null );
191
242
} else if (isRetryableError (t , maxWireVersion )) {
192
- wrapped . close ();
243
+ nullifyAndCloseWrapped ();
193
244
retryOperation (asyncBlock , callback , tryNext );
194
245
} else {
195
246
callback .onResult (null , t );
@@ -214,9 +265,15 @@ public void onResult(final AsyncBatchCursor<T> result, final Throwable t) {
214
265
if (t != null ) {
215
266
callback .onResult (null , t );
216
267
} else {
217
- wrapped = ((AsyncChangeStreamBatchCursor <T >) result ).getWrapped ();
218
- binding .release (); // release the new change stream batch cursor's reference to the binding
219
- resumeableOperation (asyncBlock , callback , tryNext );
268
+ try {
269
+ setWrappedOrCloseIt (((AsyncChangeStreamBatchCursor <T >) result ).getWrapped ());
270
+ } finally {
271
+ try {
272
+ binding .release (); // release the new change stream batch cursor's reference to the binding
273
+ } finally {
274
+ resumeableOperation (asyncBlock , callback , tryNext );
275
+ }
276
+ }
220
277
}
221
278
}
222
279
});
0 commit comments