Move error handling and span lifecycle (span.error(), span.end()) from Reactor's doOnError/doFinally operators into the async callback, before emitting the result to the subscriber.

Author: nhachicha

driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/internal/OperationExecutorImpl.java

@@ -108,20 +108,21 @@ public <T> Mono<T> execute(final ReadOperation<?, T> operation, final ReadPrefer
                                     try {
                                         binding.release();
                                     } finally {
+                                        if (t != null) {
+                                            Throwable exceptionToHandle = t instanceof MongoException
+                                                    ? OperationHelper.unwrap((MongoException) t) : t;
+                                            labelException(session, exceptionToHandle);
+                                            unpinServerAddressOnTransientTransactionError(session, exceptionToHandle);
+                                            if (span != null) {
+                                                span.error(t);
+                                            }
+                                        }
+                                        if (span != null) {
+                                            span.end();
+                                        }
                                         sinkToCallback(sink).onResult(result, t);
                                     }
-                                })).doOnError((t) -> {
-                                    Throwable exceptionToHandle = t instanceof MongoException ? OperationHelper.unwrap((MongoException) t) : t;
-                                    labelException(session, exceptionToHandle);
-                                    unpinServerAddressOnTransientTransactionError(session, exceptionToHandle);
-                                    if (span != null) {
-                                        span.error(t);
-                                    }
-                                }).doFinally(signal -> {
-                                    if (span != null) {
-                                        span.end();
-                                    }
-                                });
+                                }));
                             }
                         }).subscribe(subscriber)
         );
@@ -153,20 +154,21 @@ public <T> Mono<T> execute(final WriteOperation<T> operation, final ReadConcern
                                     try {
                                         binding.release();
                                     } finally {
+                                        if (t != null) {
+                                            Throwable exceptionToHandle = t instanceof MongoException
+                                                    ? OperationHelper.unwrap((MongoException) t) : t;
+                                            labelException(session, exceptionToHandle);
+                                            unpinServerAddressOnTransientTransactionError(session, exceptionToHandle);
+                                            if (span != null) {
+                                                span.error(t);
+                                            }
+                                        }
+                                        if (span != null) {
+                                            span.end();
+                                        }
                                         sinkToCallback(sink).onResult(result, t);
                                     }
-                                })).doOnError((t) -> {
-                                    Throwable exceptionToHandle = t instanceof MongoException ? OperationHelper.unwrap((MongoException) t) : t;
-                                    labelException(session, exceptionToHandle);
-                                    unpinServerAddressOnTransientTransactionError(session, exceptionToHandle);
-                                    if (span != null) {
-                                        span.error(t);
-                                    }
-                                }).doFinally(signal -> {
-                                    if (span != null) {
-                                        span.end();
-                                    }
-                                });
+                                }));
                                 }
                         ).subscribe(subscriber)
         );