Fixed data race in ProxyStream

Author: alex268

core/src/main/java/tech/ydb/core/utils/UpdatableOptional.java

@@ -0,0 +1,24 @@
+package tech.ydb.core.utils;
+
+import java.util.Optional;
+
+/**
+ *
+ * @author Aleksandr Gorshenin
+ * @param <R> type of value
+ */
+public class UpdatableOptional<R> {
+    private volatile R value = null;
+
+    public void update(R value) {
+        this.value = value;
+    }
+
+    public R orElse(R other) {
+        return Optional.ofNullable(value).orElse(other);
+    }
+
+    public R get() {
+        return value;
+    }
+}

query/src/main/java/tech/ydb/query/impl/SessionImpl.java

@@ -24,6 +24,7 @@
 import tech.ydb.core.operation.StatusExtractor;
 import tech.ydb.core.settings.BaseRequestSettings;
 import tech.ydb.core.utils.URITools;
+import tech.ydb.core.utils.UpdatableOptional;
 import tech.ydb.proto.query.YdbQuery;
 import tech.ydb.query.QuerySession;
 import tech.ydb.query.QueryStream;
@@ -287,9 +288,9 @@ void handleCompletion(Status status, Throwable th) { }
 
         @Override
         public CompletableFuture<Result<QueryInfo>> execute(PartsHandler handler) {
-            final CompletableFuture<Result<QueryInfo>> result = new CompletableFuture<>();
-            final AtomicReference<QueryStats> stats = new AtomicReference<>();
-            grpcStream.start(msg -> {
+            final UpdatableOptional<Status> operationStatus = new UpdatableOptional<>();
+            final UpdatableOptional<QueryStats> stats = new UpdatableOptional<>();
+            return grpcStream.start(msg -> {
                 if (isTraceEnabled) {
                     logger.trace("{} got stream message {}", SessionImpl.this, TextFormat.shortDebugString(msg));
                 }
@@ -300,7 +301,7 @@ public CompletableFuture<Result<QueryInfo>> execute(PartsHandler handler) {
 
                 if (!status.isSuccess()) {
                     handleTxMeta(null);
-                    result.complete(Result.fail(status));
+                    operationStatus.update(status);
                     return;
                 }
 
@@ -315,10 +316,7 @@ public CompletableFuture<Result<QueryInfo>> execute(PartsHandler handler) {
                     }
                 }
                 if (msg.hasExecStats()) {
-                    QueryStats old = stats.getAndSet(new QueryStats(msg.getExecStats()));
-                    if (old != null) {
-                        logger.warn("{} lost previous exec stats {}", SessionImpl.this, old);
-                    }
+                    stats.update(new QueryStats(msg.getExecStats()));
                 }
 
                 if (msg.hasResultSet()) {
@@ -329,21 +327,15 @@ public CompletableFuture<Result<QueryInfo>> execute(PartsHandler handler) {
                         logger.trace("{} lost result set part with index {}", SessionImpl.this, index);
                     }
                 }
-            }).whenComplete((status, th) -> {
-                handleCompletion(status, th);
-                if (th != null) {
-                    result.completeExceptionally(th);
-                }
-                if (status != null) {
-                    updateSessionState(status);
-                    if (status.isSuccess()) {
-                        result.complete(Result.success(new QueryInfo(stats.get()), status));
-                    } else {
-                        result.complete(Result.fail(status));
-                    }
+            }).whenComplete(this::handleCompletion).thenApply(streamStatus -> {
+                updateSessionState(streamStatus);
+                Status status = operationStatus.orElse(streamStatus);
+                if (status.isSuccess()) {
+                    return Result.success(new QueryInfo(stats.get()), streamStatus);
+                } else {
+                    return Result.fail(status);
                 }
             });
-            return result;
         }
 
         @Override

table/src/main/java/tech/ydb/table/impl/BaseSession.java

@@ -36,6 +36,7 @@
 import tech.ydb.core.settings.BaseRequestSettings;
 import tech.ydb.core.utils.ProtobufUtils;
 import tech.ydb.core.utils.URITools;
+import tech.ydb.core.utils.UpdatableOptional;
 import tech.ydb.proto.StatusCodesProtos.StatusIds;
 import tech.ydb.proto.ValueProtos;
 import tech.ydb.proto.ValueProtos.TypedValue;
@@ -1387,42 +1388,47 @@ public String toString() {
 
     private abstract class ProxyStream<R, T> implements GrpcReadStream<T> {
         private final GrpcReadStream<R> origin;
-        private final CompletableFuture<Status> result = new CompletableFuture<>();
+        private final CompletableFuture<Status> result;
+        private final UpdatableOptional<Status> operationStatus = new UpdatableOptional<>();
+        private final UpdatableOptional<Throwable> operationError = new UpdatableOptional<>();
 
         ProxyStream(GrpcReadStream<R> origin) {
             this.origin = origin;
+            this.result = new CompletableFuture<>();
         }
 
         abstract StatusIds.StatusCode readStatusCode(R message);
         abstract List<YdbIssueMessage.IssueMessage> readIssues(R message);
         abstract T readValue(R message);
 
-        private void onClose(Status status, Throwable th) {
+        private void onClose(Status streamStatus, Throwable streamError) {
+            Throwable th = operationError.orElse(streamError);
             if (th != null) {
                 updateSessionState(th, null, false);
                 result.completeExceptionally(th);
+                return;
             }
-            if (status != null) {
-                updateSessionState(null, status.getCode(), false);
-                result.complete(status);
+
+            Status st = operationStatus.orElse(streamStatus);
+            if (st != null) {
+                updateSessionState(null, st.getCode(), false);
+                result.complete(st);
             }
         }
 
         @Override
         public CompletableFuture<Status> start(Observer<T> observer) {
-            origin.start(message -> {
-                StatusIds.StatusCode statusCode = readStatusCode(message);
+            origin.start(msg -> {
+                StatusIds.StatusCode statusCode = readStatusCode(msg);
                 if (statusCode == StatusIds.StatusCode.SUCCESS) {
                     try {
-                        observer.onNext(readValue(message));
-                    } catch (Throwable t) {
-                        result.completeExceptionally(t);
+                        observer.onNext(readValue(msg));
+                    } catch (Throwable th) {
+                        operationError.update(th);
                         origin.cancel();
                     }
                 } else {
-                    Issue[] issues = Issue.fromPb(readIssues(message));
-                    StatusCode code = StatusCode.fromProto(statusCode);
-                    result.complete(Status.of(code, issues));
+                    operationStatus.update(Status.of(StatusCode.fromProto(statusCode), Issue.fromPb(readIssues(msg))));
                     origin.cancel();
                 }
             }).whenComplete(this::onClose);