// New returns an RelayService backed by a gRPC client connection. It is the
// responsibility of the caller to dial, and later close, the connection.
func New(conn *grpc.ClientConn, tracer stdopentracing.Tracer, logger log.Logger) relay.Service {
// We construct a single ratelimiter middleware, to limit the total outgoing
// QPS from this client to all methods on the remote instance. We also
// construct per-endpoint circuitbreaker middlewares to demonstrate how
// that's done, although they could easily be combined into a single breaker
// for the entire remote instance, too.
limiter := ratelimit.NewTokenBucketLimiter(jujuratelimit.NewBucketWithRate(100, 100))
var connectEndpoint endpoint.Endpoint
{
connectEndpoint = grpctransport.NewClient(
conn,
"Relay",
"Connect",
relay.EncodeGRPCConnectRequest,
relay.DecodeGRPCConnectResponse,
pb.ConnectReply{},
grpctransport.ClientBefore(opentracing.ToGRPCRequest(tracer, logger)),
).Endpoint()
connectEndpoint = opentracing.TraceClient(tracer, "Connect")(connectEndpoint)
connectEndpoint = limiter(connectEndpoint)
connectEndpoint = circuitbreaker.Gobreaker(gobreaker.NewCircuitBreaker(gobreaker.Settings{
Name: "Connect",
Timeout: 30 * time.Second,
}))(connectEndpoint)
}
var joinEndpoint endpoint.Endpoint
{
joinEndpoint = grpctransport.NewClient(
conn,
"Relay",
"Join",
relay.EncodeGRPCJoinRequest,
relay.DecodeGRPCJoinResponse,
pb.JoinReply{},
grpctransport.ClientBefore(opentracing.ToGRPCRequest(tracer, logger)),
).Endpoint()
joinEndpoint = opentracing.TraceClient(tracer, "Join")(joinEndpoint)
joinEndpoint = limiter(joinEndpoint)
joinEndpoint = circuitbreaker.Gobreaker(gobreaker.NewCircuitBreaker(gobreaker.Settings{
Name: "Join",
Timeout: 30 * time.Second,
}))(joinEndpoint)
}
return relay.Endpoints{
ConnectEndpoint: connectEndpoint,
JoinEndpoint: joinEndpoint,
}
}
func TestTraceGRPCRequestRoundtrip(t *testing.T) {
logger := log.NewNopLogger()
tracer := mocktracer.New()
// Initialize the ctx with a Span to inject.
beforeSpan := tracer.StartSpan("to_inject").(*mocktracer.MockSpan)
defer beforeSpan.Finish()
beforeSpan.Context().SetBaggageItem("baggage", "check")
beforeCtx := opentracing.ContextWithSpan(context.Background(), beforeSpan)
toGRPCFunc := kitot.ToGRPCRequest(tracer, logger)
md := metadata.Pairs()
// Call the RequestFunc.
afterCtx := toGRPCFunc(beforeCtx, &md)
// The Span should not have changed.
afterSpan := opentracing.SpanFromContext(afterCtx)
if beforeSpan != afterSpan {
t.Errorf("Should not swap in a new span")
}
// No spans should have finished yet.
finishedSpans := tracer.GetFinishedSpans()
if want, have := 0, len(finishedSpans); want != have {
t.Errorf("Want %v span(s), found %v", want, have)
}
// Use FromGRPCRequest to verify that we can join with the trace given MD.
fromGRPCFunc := kitot.FromGRPCRequest(tracer, "joined", logger)
joinCtx := fromGRPCFunc(afterCtx, &md)
joinedSpan := opentracing.SpanFromContext(joinCtx).(*mocktracer.MockSpan)
joinedContext := joinedSpan.Context().(*mocktracer.MockSpanContext)
beforeContext := beforeSpan.Context().(*mocktracer.MockSpanContext)
if joinedContext.SpanID == beforeContext.SpanID {
t.Error("SpanID should have changed", joinedContext.SpanID, beforeContext.SpanID)
}
// Check that the parent/child relationship is as expected for the joined span.
if want, have := beforeContext.SpanID, joinedSpan.ParentID; want != have {
t.Errorf("Want ParentID %q, have %q", want, have)
}
if want, have := "joined", joinedSpan.OperationName; want != have {
t.Errorf("Want %q, have %q", want, have)
}
if want, have := "check", joinedSpan.Context().BaggageItem("baggage"); want != have {
t.Errorf("Want %q, have %q", want, have)
}
}
// MakeConcatEndpointFactory returns a loadbalancer.Factory that transforms
// GRPC host:port strings into Endpoints that call the Concat method on a GRPC
// server at that address.
func MakeConcatEndpointFactory(tracer opentracing.Tracer, tracingLogger log.Logger) loadbalancer.Factory {
return func(instance string) (endpoint.Endpoint, io.Closer, error) {
cc, err := grpc.Dial(instance, grpc.WithInsecure())
return grpctransport.NewClient(
cc,
"Add",
"Concat",
encodeConcatRequest,
decodeConcatResponse,
pb.ConcatReply{},
grpctransport.SetClientBefore(kitot.ToGRPCRequest(tracer, tracingLogger)),
).Endpoint(), cc, err
}
}