use ServerStreaming scenario to demonstrate graceful shutdown

Author: purnesh42H

examples/features/gracefulstop/README.md

@@ -1,31 +1,33 @@
 # Graceful Stop
 
 This example demonstrates how to gracefully stop a gRPC server using
-Server.GracefulStop().
+`Server.GracefulStop()`.
 
 ## How to run
 
-Start the server which will listen to incoming gRPC requests as well as OS
-interrupt signals (SIGINT/SIGTERM). After receiving interrupt signal, it calls
-`s.GracefulStop()` to gracefully shut down the server. If graceful shutdown
-doesn't happen in time, server is stopped forcefully.
+Start the server which will serve `ServerStreaming` gRPC requests as well as an
+OS interrupt signal (SIGINT/SIGTERM). `ServerStreaming` handler returns an
+error after sending 5 messages to the client or if client's context is
+canceled. If an interrupt signal is received, it calls `s.GracefulStop()` to
+gracefully shut down the server. If graceful shutdown doesn't happen in time,
+the server is stopped forcefully.
 
 ```sh
 $ go run server/main.go
 ```
 
-In a separate terminal, start the client which will send multiple requests to
-the server with some delay between each request.
+In a separate terminal, start the client which will start `ServerStreaming`
+request to the server and keep receiving messages until an error is received.
+Once an error is received, it closes the stream.
 
 ```sh
 $ go run client/main.go
 ```
 
-Use Ctrl+C or SIGTERM to signal the server to shut down.
+Once the client starts receiving messages from server, use Ctrl+C or SIGTERM to
+signal the server to shut down.
 
-The server begins a graceful stop:
-- It finishes ongoing requests.
-- Rejects new incoming requests.
-
-The client will notice the server's shutdown when a request fails.
+The server begins a graceful stop. It finish the in-flight request. In this
+case, client will receive 5 messages followed by the stream failure from the
+server, allowing client to close the stream gracefully, before shutting down.
 

examples/features/gracefulstop/client/main.go

@@ -26,41 +26,43 @@ import (
 	"time"
 
 	"google.golang.org/grpc"
-	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials/insecure"
-	pb "google.golang.org/grpc/examples/helloworld/helloworld"
-	"google.golang.org/grpc/status"
+	pb "google.golang.org/grpc/examples/features/proto/echo"
 )
 
 var addr = flag.String("addr", "localhost:50052", "the address to connect to")
 
 func main() {
 	flag.Parse()
 
-	// Set up a connection to the server.
 	conn, err := grpc.NewClient(*addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
 	if err != nil {
-		log.Fatalf("Failed to connect: %v", err)
+		log.Fatalf("Failed to create new client: %v", err)
 	}
 	defer conn.Close()
-	c := pb.NewGreeterClient(conn)
+	c := pb.NewEchoClient(conn)
 
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
 	defer cancel()
 
-	for i := 1; i <= 10; i++ {
-		log.Printf("Calling SayHello %d time", i)
-		r, err := c.SayHello(ctx, &pb.HelloRequest{})
+	stream, err := c.ServerStreamingEcho(ctx, &pb.EchoRequest{})
+	if err != nil {
+		log.Fatalf("Error starting stream: %v", err)
+	}
+
+	for {
+		r, err := stream.Recv()
 		if err != nil {
-			if status.Code(err) != codes.InvalidArgument {
-				log.Printf("Received unexpected error: %v", err)
-				continue
+			// Handle the error and close the stream gracefully
+			log.Printf("Error sending request: %v\n", err)
+			err := stream.CloseSend()
+			if err != nil {
+				log.Fatalf("Error closing stream: %v", err)
 			}
-			log.Printf("Received error: %v", err)
-			continue
+			log.Println("Stream closed gracefully")
+			break
 		}
-		log.Printf("Received response: %s", r.Message)
-		time.Sleep(time.Second)
+		log.Printf(r.Message)
 	}
 
 	log.Printf("Client finished interaction with server.")

examples/features/gracefulstop/server/main.go

@@ -60,15 +60,12 @@ func (s *server) ServerStreamingEcho(_ *pb.EchoRequest, stream pb.Echo_ServerStr
 			atomic.AddInt32(&streamMessages, 1)
 
 			mu.Lock()
-
 			if streamMessages > 5 {
 				return fmt.Errorf("request failed")
 			}
-
 			if err := stream.Send(&pb.EchoResponse{Message: fmt.Sprintf("Messages Sent: %d", streamMessages)}); err != nil {
 				return err
 			}
-
 			mu.Unlock()
 
 			time.Sleep(1 * time.Second)