Merge pull request #14080 from yang-g/return_early

Clear pending data when stream is closed by server.

5 files changed, 251 insertions, 39 deletions

diff --git a/test/cpp/end2end/BUILD b/test/cpp/end2end/BUILD
index 8894c68b95..27e8da1ff1 100644
--- a/test/cpp/end2end/BUILD
+++ b/test/cpp/end2end/BUILD
@@ -120,6 +120,24 @@ grpc_cc_library(
 )
 
 grpc_cc_test(
+    name = "server_early_return_test",
+    srcs = ["server_early_return_test.cc"],
+    external_deps = [
+        "gtest",
+    ],
+    deps = [
+        "//:gpr",
+        "//:grpc",
+        "//:grpc++",
+        "//src/proto/grpc/testing:echo_messages_proto",
+        "//src/proto/grpc/testing:echo_proto",
+        "//test/core/util:gpr_test_util",
+        "//test/core/util:grpc_test_util",
+        "//test/cpp/util:test_util",
+    ],
+)
+
+grpc_cc_test(
     name = "end2end_test",
     deps = [
         ":end2end_test_lib",
diff --git a/test/cpp/end2end/end2end_test.cc b/test/cpp/end2end/end2end_test.cc
index e25de810f1..b1b0b5d0e8 100644
--- a/test/cpp/end2end/end2end_test.cc
+++ b/test/cpp/end2end/end2end_test.cc
@@ -1087,31 +1087,6 @@ TEST_P(End2endTest, RpcMaxMessageSize) {
   EXPECT_FALSE(s.ok());
 }
 
-// Client sends 20 requests and the server returns CANCELLED status after
-// reading 10 requests.
-TEST_P(End2endTest, RequestStreamServerEarlyCancelTest) {
-  ResetStub();
-  EchoRequest request;
-  EchoResponse response;
-  ClientContext context;
-
-  context.AddMetadata(kServerCancelAfterReads, "10");
-  auto stream = stub_->RequestStream(&context, &response);
-  request.set_message("hello");
-  int send_messages = 20;
-  while (send_messages > 10) {
-    EXPECT_TRUE(stream->Write(request));
-    send_messages--;
-  }
-  while (send_messages > 0) {
-    stream->Write(request);
-    send_messages--;
-  }
-  stream->WritesDone();
-  Status s = stream->Finish();
-  EXPECT_EQ(s.error_code(), StatusCode::CANCELLED);
-}
-
 void ReaderThreadFunc(ClientReaderWriter<EchoRequest, EchoResponse>* stream,
                       gpr_event* ev) {
   EchoResponse resp;
diff --git a/test/cpp/end2end/server_early_return_test.cc b/test/cpp/end2end/server_early_return_test.cc
new file mode 100644
index 0000000000..6fd26fcdb6
--- /dev/null
+++ b/test/cpp/end2end/server_early_return_test.cc
@@ -0,0 +1,233 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+#include <grpc++/channel.h>
+#include <grpc++/client_context.h>
+#include <grpc++/create_channel.h>
+#include <grpc++/security/credentials.h>
+#include <grpc++/security/server_credentials.h>
+#include <grpc++/server.h>
+#include <grpc++/server_builder.h>
+#include <grpc++/server_context.h>
+#include <grpc/grpc.h>
+#include <grpc/support/alloc.h>
+#include <grpc/support/log.h>
+#include <grpc/support/thd.h>
+#include <grpc/support/time.h>
+
+#include "src/proto/grpc/testing/echo.grpc.pb.h"
+#include "test/core/util/port.h"
+#include "test/core/util/test_config.h"
+#include "test/cpp/util/string_ref_helper.h"
+
+#include <gtest/gtest.h>
+
+namespace grpc {
+namespace testing {
+namespace {
+
+const char kServerReturnStatusCode[] = "server_return_status_code";
+const char kServerDelayBeforeReturnUs[] = "server_delay_before_return_us";
+const char kServerReturnAfterNReads[] = "server_return_after_n_reads";
+
+class TestServiceImpl : public ::grpc::testing::EchoTestService::Service {
+ public:
+  // Unused methods are not implemented.
+
+  Status RequestStream(ServerContext* context,
+                       ServerReader<EchoRequest>* reader,
+                       EchoResponse* response) override {
+    int server_return_status_code =
+        GetIntValueFromMetadata(context, kServerReturnStatusCode, 0);
+    int server_delay_before_return_us =
+        GetIntValueFromMetadata(context, kServerDelayBeforeReturnUs, 0);
+    int server_return_after_n_reads =
+        GetIntValueFromMetadata(context, kServerReturnAfterNReads, 0);
+
+    EchoRequest request;
+    while (server_return_after_n_reads--) {
+      EXPECT_TRUE(reader->Read(&request));
+    }
+
+    response->set_message("response msg");
+
+    gpr_sleep_until(gpr_time_add(
+        gpr_now(GPR_CLOCK_MONOTONIC),
+        gpr_time_from_micros(server_delay_before_return_us, GPR_TIMESPAN)));
+
+    return Status(static_cast<StatusCode>(server_return_status_code), "");
+  }
+
+  Status BidiStream(
+      ServerContext* context,
+      ServerReaderWriter<EchoResponse, EchoRequest>* stream) override {
+    int server_return_status_code =
+        GetIntValueFromMetadata(context, kServerReturnStatusCode, 0);
+    int server_delay_before_return_us =
+        GetIntValueFromMetadata(context, kServerDelayBeforeReturnUs, 0);
+    int server_return_after_n_reads =
+        GetIntValueFromMetadata(context, kServerReturnAfterNReads, 0);
+
+    EchoRequest request;
+    EchoResponse response;
+    while (server_return_after_n_reads--) {
+      EXPECT_TRUE(stream->Read(&request));
+      response.set_message(request.message());
+      EXPECT_TRUE(stream->Write(response));
+    }
+
+    gpr_sleep_until(gpr_time_add(
+        gpr_now(GPR_CLOCK_MONOTONIC),
+        gpr_time_from_micros(server_delay_before_return_us, GPR_TIMESPAN)));
+
+    return Status(static_cast<StatusCode>(server_return_status_code), "");
+  }
+
+  int GetIntValueFromMetadata(ServerContext* context, const char* key,
+                              int default_value) {
+    auto metadata = context->client_metadata();
+    if (metadata.find(key) != metadata.end()) {
+      std::istringstream iss(ToString(metadata.find(key)->second));
+      iss >> default_value;
+    }
+    return default_value;
+  }
+};
+
+class ServerEarlyReturnTest : public ::testing::Test {
+ protected:
+  ServerEarlyReturnTest() : picked_port_(0) {}
+
+  void SetUp() override {
+    int port = grpc_pick_unused_port_or_die();
+    picked_port_ = port;
+    server_address_ << "127.0.0.1:" << port;
+    ServerBuilder builder;
+    builder.AddListeningPort(server_address_.str(),
+                             InsecureServerCredentials());
+    builder.RegisterService(&service_);
+    server_ = builder.BuildAndStart();
+
+    channel_ =
+        CreateChannel(server_address_.str(), InsecureChannelCredentials());
+    stub_ = grpc::testing::EchoTestService::NewStub(channel_);
+  }
+
+  void TearDown() override {
+    server_->Shutdown();
+    if (picked_port_ > 0) {
+      grpc_recycle_unused_port(picked_port_);
+    }
+  }
+
+  // Client sends 20 requests and the server returns after reading 10 requests.
+  // If return_cancel is true, server returns CANCELLED status. Otherwise it
+  // returns OK.
+  void DoBidiStream(bool return_cancelled) {
+    EchoRequest request;
+    EchoResponse response;
+    ClientContext context;
+
+    context.AddMetadata(kServerReturnAfterNReads, "10");
+    if (return_cancelled) {
+      // "1" means CANCELLED
+      context.AddMetadata(kServerReturnStatusCode, "1");
+    }
+    context.AddMetadata(kServerDelayBeforeReturnUs, "10000");
+
+    auto stream = stub_->BidiStream(&context);
+
+    for (int i = 0; i < 20; i++) {
+      request.set_message(grpc::string("hello") + grpc::to_string(i));
+      bool write_ok = stream->Write(request);
+      bool read_ok = stream->Read(&response);
+      if (i < 10) {
+        EXPECT_TRUE(write_ok);
+        EXPECT_TRUE(read_ok);
+        EXPECT_EQ(response.message(), request.message());
+      } else {
+        EXPECT_FALSE(read_ok);
+      }
+    }
+
+    stream->WritesDone();
+    EXPECT_FALSE(stream->Read(&response));
+
+    Status s = stream->Finish();
+    if (return_cancelled) {
+      EXPECT_EQ(s.error_code(), StatusCode::CANCELLED);
+    } else {
+      EXPECT_TRUE(s.ok());
+    }
+  }
+
+  void DoRequestStream(bool return_cancelled) {
+    EchoRequest request;
+    EchoResponse response;
+    ClientContext context;
+
+    context.AddMetadata(kServerReturnAfterNReads, "10");
+    if (return_cancelled) {
+      // "1" means CANCELLED
+      context.AddMetadata(kServerReturnStatusCode, "1");
+    }
+    context.AddMetadata(kServerDelayBeforeReturnUs, "10000");
+
+    auto stream = stub_->RequestStream(&context, &response);
+    for (int i = 0; i < 20; i++) {
+      request.set_message(grpc::string("hello") + grpc::to_string(i));
+      bool written = stream->Write(request);
+      if (i < 10) {
+        EXPECT_TRUE(written);
+      }
+    }
+    stream->WritesDone();
+    Status s = stream->Finish();
+    if (return_cancelled) {
+      EXPECT_EQ(s.error_code(), StatusCode::CANCELLED);
+    } else {
+      EXPECT_TRUE(s.ok());
+    }
+  }
+
+  std::shared_ptr<Channel> channel_;
+  std::unique_ptr<grpc::testing::EchoTestService::Stub> stub_;
+  std::unique_ptr<Server> server_;
+  std::ostringstream server_address_;
+  TestServiceImpl service_;
+  int picked_port_;
+};
+
+TEST_F(ServerEarlyReturnTest, BidiStreamEarlyOk) { DoBidiStream(false); }
+
+TEST_F(ServerEarlyReturnTest, BidiStreamEarlyCancel) { DoBidiStream(true); }
+
+TEST_F(ServerEarlyReturnTest, RequestStreamEarlyOK) { DoRequestStream(false); }
+TEST_F(ServerEarlyReturnTest, RequestStreamEarlyCancel) {
+  DoRequestStream(true);
+}
+
+}  // namespace
+}  // namespace testing
+}  // namespace grpc
+
+int main(int argc, char** argv) {
+  grpc_test_init(argc, argv);
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}
diff --git a/test/cpp/end2end/test_service_impl.cc b/test/cpp/end2end/test_service_impl.cc
index e4f7c08f25..875a27b5f5 100644
--- a/test/cpp/end2end/test_service_impl.cc
+++ b/test/cpp/end2end/test_service_impl.cc
@@ -181,13 +181,6 @@ Status TestServiceImpl::RequestStream(ServerContext* context,
   int server_try_cancel = GetIntValueFromMetadata(
       kServerTryCancelRequest, context->client_metadata(), DO_NOT_CANCEL);
 
-  // If 'cancel_after_reads' is set in the metadata AND non-zero, the server
-  // will cancel the RPC (by just returning Status::CANCELLED - doesn't call
-  // ServerContext::TryCancel()) after reading the number of records specified
-  // by the 'cancel_after_reads' value set in the metadata.
-  int cancel_after_reads = GetIntValueFromMetadata(
-      kServerCancelAfterReads, context->client_metadata(), 0);
-
   EchoRequest request;
   response->set_message("");
 
@@ -204,12 +197,6 @@ Status TestServiceImpl::RequestStream(ServerContext* context,
 
   int num_msgs_read = 0;
   while (reader->Read(&request)) {
-    if (cancel_after_reads == 1) {
-      gpr_log(GPR_INFO, "return cancel status");
-      return Status::CANCELLED;
-    } else if (cancel_after_reads > 0) {
-      cancel_after_reads--;
-    }
     response->mutable_message()->append(request.message());
   }
   gpr_log(GPR_INFO, "Read: %d messages", num_msgs_read);
diff --git a/test/cpp/end2end/test_service_impl.h b/test/cpp/end2end/test_service_impl.h
index e485769bb2..070f357139 100644
--- a/test/cpp/end2end/test_service_impl.h
+++ b/test/cpp/end2end/test_service_impl.h
@@ -31,7 +31,6 @@ namespace testing {
 
 const int kServerDefaultResponseStreamsToSend = 3;
 const char* const kServerResponseStreamsToSend = "server_responses_to_send";
-const char* const kServerCancelAfterReads = "cancel_after_reads";
 const char* const kServerTryCancelRequest = "server_try_cancel";
 const char* const kDebugInfoTrailerKey = "debug-info-bin";
 const char* const kServerFinishAfterNReads = "server_finish_after_n_reads";