github.com/newrelic/go-agent@v3.26.0+incompatible/_integrations/nrgrpc/testapp/server.go (about)

     1  // Copyright 2020 New Relic Corporation. All rights reserved.
     2  // SPDX-License-Identifier: Apache-2.0
     3  
     4  package testapp
     5  
     6  import (
     7  	"context"
     8  	"encoding/json"
     9  	"io"
    10  
    11  	newrelic "github.com/newrelic/go-agent"
    12  	codes "google.golang.org/grpc/codes"
    13  	"google.golang.org/grpc/metadata"
    14  	status "google.golang.org/grpc/status"
    15  )
    16  
    17  // Server is a gRPC server.
    18  type Server struct{}
    19  
    20  // DoUnaryUnary is a unary request, unary response method.
    21  func (s *Server) DoUnaryUnary(ctx context.Context, msg *Message) (*Message, error) {
    22  	defer newrelic.StartSegment(newrelic.FromContext(ctx), "DoUnaryUnary").End()
    23  	md, _ := metadata.FromIncomingContext(ctx)
    24  	js, _ := json.Marshal(md)
    25  	return &Message{Text: string(js)}, nil
    26  }
    27  
    28  // DoUnaryStream is a unary request, stream response method.
    29  func (s *Server) DoUnaryStream(msg *Message, stream TestApplication_DoUnaryStreamServer) error {
    30  	defer newrelic.StartSegment(newrelic.FromContext(stream.Context()), "DoUnaryStream").End()
    31  	md, _ := metadata.FromIncomingContext(stream.Context())
    32  	js, _ := json.Marshal(md)
    33  	for i := 0; i < 3; i++ {
    34  		if err := stream.Send(&Message{Text: string(js)}); nil != err {
    35  			return err
    36  		}
    37  	}
    38  	return nil
    39  }
    40  
    41  // DoStreamUnary is a stream request, unary response method.
    42  func (s *Server) DoStreamUnary(stream TestApplication_DoStreamUnaryServer) error {
    43  	defer newrelic.StartSegment(newrelic.FromContext(stream.Context()), "DoStreamUnary").End()
    44  	md, _ := metadata.FromIncomingContext(stream.Context())
    45  	js, _ := json.Marshal(md)
    46  	for {
    47  		_, err := stream.Recv()
    48  		if err == io.EOF {
    49  			return stream.SendAndClose(&Message{Text: string(js)})
    50  		} else if nil != err {
    51  			return err
    52  		}
    53  	}
    54  }
    55  
    56  // DoStreamStream is a stream request, stream response method.
    57  func (s *Server) DoStreamStream(stream TestApplication_DoStreamStreamServer) error {
    58  	defer newrelic.StartSegment(newrelic.FromContext(stream.Context()), "DoStreamStream").End()
    59  	md, _ := metadata.FromIncomingContext(stream.Context())
    60  	js, _ := json.Marshal(md)
    61  	for {
    62  		_, err := stream.Recv()
    63  		if err == io.EOF {
    64  			return nil
    65  		} else if nil != err {
    66  			return err
    67  		}
    68  		if err := stream.Send(&Message{Text: string(js)}); nil != err {
    69  			return err
    70  		}
    71  	}
    72  }
    73  
    74  // DoUnaryUnaryError is a unary request, unary response method that returns an
    75  // error.
    76  func (s *Server) DoUnaryUnaryError(ctx context.Context, msg *Message) (*Message, error) {
    77  	return &Message{}, status.New(codes.DataLoss, "oooooops!").Err()
    78  }
    79  
    80  // DoUnaryStreamError is a unary request, unary response method that returns an
    81  // error.
    82  func (s *Server) DoUnaryStreamError(msg *Message, stream TestApplication_DoUnaryStreamErrorServer) error {
    83  	return status.New(codes.DataLoss, "oooooops!").Err()
    84  }