Added new implementatin of grpc stream proxies

Author: alex268

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

@@ -20,7 +20,6 @@
 import tech.ydb.core.StatusCode;
 import tech.ydb.core.grpc.GrpcReadStream;
 import tech.ydb.core.grpc.GrpcRequestSettings;
-import tech.ydb.core.impl.call.ProxyReadStream;
 import tech.ydb.core.operation.StatusExtractor;
 import tech.ydb.core.settings.BaseRequestSettings;
 import tech.ydb.core.utils.URITools;
@@ -137,12 +136,27 @@ GrpcReadStream<Status> attach(AttachSessionSettings settings) {
                 .setSessionId(sessionId)
                 .build();
         GrpcRequestSettings grpcSettings = makeOptions(settings).build();
-        return new ProxyReadStream<>(rpc.attachSession(request, grpcSettings), (message, promise, observer) -> {
-            logger.trace("session '{}' got attach stream message {}", sessionId, TextFormat.shortDebugString(message));
-            Status status = Status.of(StatusCode.fromProto(message.getStatus()), Issue.fromPb(message.getIssuesList()));
-            updateSessionState(status);
-            observer.onNext(status);
-        });
+        GrpcReadStream<YdbQuery.SessionState> origin = rpc.attachSession(request, grpcSettings);
+        return new GrpcReadStream<Status>() {
+            @Override
+            public CompletableFuture<Status> start(GrpcReadStream.Observer<Status> observer) {
+                return origin.start(message -> {
+                    if (logger.isTraceEnabled()) {
+                        String msg = TextFormat.shortDebugString(message);
+                        logger.trace("session '{}' got attach stream message {}", sessionId, msg);
+                    }
+                    StatusCode code = StatusCode.fromProto(message.getStatus());
+                    Status status = Status.of(code, Issue.fromPb(message.getIssuesList()));
+                    updateSessionState(status);
+                    observer.onNext(status);
+                });
+            }
+
+            @Override
+            public void cancel() {
+                origin.cancel();
+            }
+        };
     }
 
     private GrpcRequestSettings.Builder makeOptions(BaseRequestSettings settings) {

table/src/main/java/tech/ydb/table/Session.java

@@ -9,7 +9,6 @@
 import tech.ydb.core.Result;
 import tech.ydb.core.Status;
 import tech.ydb.core.grpc.GrpcReadStream;
-import tech.ydb.core.impl.call.ProxyReadStream;
 import tech.ydb.table.description.TableDescription;
 import tech.ydb.table.description.TableOptionDescription;
 import tech.ydb.table.query.DataQuery;
@@ -151,9 +150,18 @@ default CompletableFuture<Result<TableTransaction>> beginTransaction(TxMode txMo
 
     @Deprecated
     default GrpcReadStream<ResultSetReader> readTable(String tablePath, ReadTableSettings settings) {
-        return new ProxyReadStream<>(executeReadTable(tablePath, settings), (part, promise, observer) -> {
-            observer.onNext(part.getResultSetReader());
-        });
+        GrpcReadStream<ReadTablePart> stream = executeReadTable(tablePath, settings);
+        return new GrpcReadStream<ResultSetReader>() {
+            @Override
+            public CompletableFuture<Status> start(GrpcReadStream.Observer<ResultSetReader> observer) {
+                return stream.start(part -> observer.onNext(part.getResultSetReader()));
+            }
+
+            @Override
+            public void cancel() {
+                stream.cancel();
+            }
+        };
     }
 
     @Deprecated

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

@@ -31,14 +31,14 @@
 import tech.ydb.core.grpc.GrpcReadStream;
 import tech.ydb.core.grpc.GrpcRequestSettings;
 import tech.ydb.core.grpc.YdbHeaders;
-import tech.ydb.core.impl.call.ProxyReadStream;
 import tech.ydb.core.operation.Operation;
 import tech.ydb.core.settings.BaseRequestSettings;
 import tech.ydb.core.utils.ProtobufUtils;
 import tech.ydb.core.utils.URITools;
 import tech.ydb.proto.StatusCodesProtos.StatusIds;
 import tech.ydb.proto.ValueProtos;
 import tech.ydb.proto.ValueProtos.TypedValue;
+import tech.ydb.proto.YdbIssueMessage;
 import tech.ydb.proto.common.CommonProtos;
 import tech.ydb.proto.scheme.SchemeOperationProtos;
 import tech.ydb.proto.table.YdbTable;
@@ -1205,22 +1205,22 @@ public GrpcReadStream<ReadTablePart> executeReadTable(String tablePath, ReadTabl
         }
 
         GrpcReadStream<YdbTable.ReadTableResponse> origin = rpc.streamReadTable(request.build(), options.build());
-        return new ProxyReadStream<>(origin, (response, future, observer) -> {
-            StatusIds.StatusCode statusCode = response.getStatus();
-            if (statusCode == StatusIds.StatusCode.SUCCESS) {
-                try {
-                    observer.onNext(new ReadTablePart(response.getResult(), response.getSnapshot()));
-                } catch (Throwable t) {
-                    future.completeExceptionally(t);
-                    origin.cancel();
-                }
-            } else {
-                Issue[] issues = Issue.fromPb(response.getIssuesList());
-                StatusCode code = StatusCode.fromProto(statusCode);
-                future.complete(Status.of(code, issues));
-                origin.cancel();
+        return new ProxyStream<YdbTable.ReadTableResponse, ReadTablePart>(origin) {
+            @Override
+            StatusIds.StatusCode readStatusCode(YdbTable.ReadTableResponse message) {
+                return message.getStatus();
             }
-        });
+
+            @Override
+            List<YdbIssueMessage.IssueMessage> readIssues(YdbTable.ReadTableResponse message) {
+                return message.getIssuesList();
+            }
+
+            @Override
+            ReadTablePart readValue(YdbTable.ReadTableResponse message) {
+                return new ReadTablePart(message.getResult(), message.getSnapshot());
+            }
+        };
     }
 
     @Override
@@ -1239,22 +1239,22 @@ public GrpcReadStream<ResultSetReader> executeScanQuery(String query, Params par
         }
 
         GrpcReadStream<YdbTable.ExecuteScanQueryPartialResponse> origin = rpc.streamExecuteScanQuery(req, opts.build());
-        return new ProxyReadStream<>(origin, (response, future, observer) -> {
-            StatusIds.StatusCode statusCode = response.getStatus();
-            if (statusCode == StatusIds.StatusCode.SUCCESS) {
-                try {
-                    observer.onNext(ProtoValueReaders.forResultSet(response.getResult().getResultSet()));
-                } catch (Throwable t) {
-                    future.completeExceptionally(t);
-                    origin.cancel();
-                }
-            } else {
-                Issue[] issues = Issue.fromPb(response.getIssuesList());
-                StatusCode code = StatusCode.fromProto(statusCode);
-                future.complete(Status.of(code, issues));
-                origin.cancel();
+        return new ProxyStream<YdbTable.ExecuteScanQueryPartialResponse, ResultSetReader>(origin) {
+            @Override
+            StatusIds.StatusCode readStatusCode(YdbTable.ExecuteScanQueryPartialResponse message) {
+                return message.getStatus();
             }
-        });
+
+            @Override
+            List<YdbIssueMessage.IssueMessage> readIssues(YdbTable.ExecuteScanQueryPartialResponse message) {
+                return message.getIssuesList();
+            }
+
+            @Override
+            ResultSetReader readValue(YdbTable.ExecuteScanQueryPartialResponse message) {
+                return ProtoValueReaders.forResultSet(message.getResult().getResultSet());
+            }
+        };
     }
 
     private CompletableFuture<Status> commitTransactionInternal(String txId, CommitTxSettings settings) {
@@ -1378,6 +1378,56 @@ public String toString() {
         return "Session{" + id + "}";
     }
 
+    private abstract class ProxyStream<R, T> implements GrpcReadStream<T> {
+        private final GrpcReadStream<R> origin;
+        private final CompletableFuture<Status> result = new CompletableFuture<>();
+
+        ProxyStream(GrpcReadStream<R> origin) {
+            this.origin = origin;
+        }
+
+        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) {
+            if (th != null) {
+                updateSessionState(th, null, false);
+                result.completeExceptionally(th);
+            }
+            if (status != null) {
+                updateSessionState(null, status.getCode(), false);
+                result.complete(status);
+            }
+        }
+
+        @Override
+        public CompletableFuture<Status> start(Observer<T> observer) {
+            origin.start(message -> {
+                StatusIds.StatusCode statusCode = readStatusCode(message);
+                if (statusCode == StatusIds.StatusCode.SUCCESS) {
+                    try {
+                        observer.onNext(readValue(message));
+                    } catch (Throwable t) {
+                        result.completeExceptionally(t);
+                        origin.cancel();
+                    }
+                } else {
+                    Issue[] issues = Issue.fromPb(readIssues(message));
+                    StatusCode code = StatusCode.fromProto(statusCode);
+                    result.complete(Status.of(code, issues));
+                    origin.cancel();
+                }
+            }).whenComplete(this::onClose);
+            return result;
+        }
+
+        @Override
+        public void cancel() {
+            origin.cancel();
+        }
+    }
+
     class TableTransactionImpl extends YdbTransactionImpl implements TableTransaction {
 
         TableTransactionImpl(TxMode txMode, String txId) {