examples/features/deadline/server/main.go

/*
 *
 * 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.
 *
 */

// Binary server demonstrates how to handle RPCs with deadlines and how to
// propagate deadlines in requests.
package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"strings"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/status"

	pb "google.golang.org/grpc/examples/features/proto/echo"
)

var port = flag.Int("port", 50052, "port number")

// server is used to implement EchoServer.
type server struct {
	pb.UnimplementedEchoServer
	client pb.EchoClient
	cc     *grpc.ClientConn
}

func (s *server) UnaryEcho(ctx context.Context, req *pb.EchoRequest) (*pb.EchoResponse, error) {
	message := req.Message
	if strings.HasPrefix(message, "[propagate me]") {
		time.Sleep(800 * time.Millisecond)
		message = strings.TrimPrefix(message, "[propagate me]")
		return s.client.UnaryEcho(ctx, &pb.EchoRequest{Message: message})
	}

	if message == "delay" {
		time.Sleep(1500 * time.Millisecond)
	}

	return &pb.EchoResponse{Message: req.Message}, nil
}

func (s *server) BidirectionalStreamingEcho(stream pb.Echo_BidirectionalStreamingEchoServer) error {
	for {
		req, err := stream.Recv()
		if err == io.EOF {
			return status.Error(codes.InvalidArgument, "request message not received")
		}
		if err != nil {
			return err
		}

		message := req.Message
		if strings.HasPrefix(message, "[propagate me]") {
			time.Sleep(800 * time.Millisecond)
			message = strings.TrimPrefix(message, "[propagate me]")
			res, err := s.client.UnaryEcho(stream.Context(), &pb.EchoRequest{Message: message})
			if err != nil {
				return err
			}
			stream.Send(res)
		}

		if message == "delay" {
			time.Sleep(1500 * time.Millisecond)
		}
		stream.Send(&pb.EchoResponse{Message: message})
	}
}

func (s *server) Close() {
	s.cc.Close()
}

func newEchoServer() *server {
	target := fmt.Sprintf("localhost:%v", *port)
	cc, err := grpc.NewClient(target, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		log.Fatalf("did not connect: %v", err)
	}
	return &server{client: pb.NewEchoClient(cc), cc: cc}
}

func main() {
	flag.Parse()

	address := fmt.Sprintf(":%v", *port)
	lis, err := net.Listen("tcp", address)
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}

	echoServer := newEchoServer()
	defer echoServer.Close()

	grpcServer := grpc.NewServer()
	pb.RegisterEchoServer(grpcServer, echoServer)

	if err := grpcServer.Serve(lis); err != nil {
		log.Fatalf("failed to serve: %v", err)
	}
}